# imager_utils.py
# General imager functions for total intensity VLBI data
#
#    Copyright (C) 2018 Andrew Chael
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.


from __future__ import division
from __future__ import print_function

from builtins import str
from builtins import range
from builtins import object

import time
import numpy as np
import scipy.optimize as opt
import matplotlib.pyplot as plt

import ehtim.image as image
import ehtim.observing.obs_helpers as obsh
import ehtim.const_def as ehc


##################################################################################################
# Constants & Definitions
##################################################################################################

NORM_REGULARIZER = False  # ANDREW TODO change this default in the future

MAXLS = 100  # maximum number of line searches in L-BFGS-B
NHIST = 100  # number of steps to store for hessian approx
MAXIT = 100  # maximum number of iterations
STOP = 1.e-8  # convergence criterion

DATATERMS = ['vis', 'bs', 'amp', 'cphase', 'cphase_diag',
             'camp', 'logcamp', 'logcamp_diag', 'logamp']
REGULARIZERS = ['gs', 'tv', 'tv2', 'l1w', 'lA', 'patch', 'simple', 'compact', 'compact2', 'rgauss']

nit = 0  # global variable to track the iteration number in the plotting callback

##################################################################################################
# Total Intensity Imager
##################################################################################################


def imager_func(Obsdata, InitIm, Prior, flux,
                d1='vis', d2=False, d3=False,
                alpha_d1=100, alpha_d2=100, alpha_d3=100,
                s1='simple', s2=False, s3=False,
                alpha_s1=1, alpha_s2=1, alpha_s3=1,
                   alpha_flux=500, alpha_cm=500,
                **kwargs):
    """Run a general interferometric imager. Only works directly on the image's primary polarization.

       Args:
           Obsdata (Obsdata): The Obsdata object with VLBI data
           InitIm (Image): The Image object with the initial image for the minimization
           Prior (Image): The Image object with the prior image
           flux (float): The total flux of the output image in Jy

           d1 (str): The first data term; options are 'vis', 'bs', 'amp', 'cphase', 'cphase_diag' 'camp', 'logcamp', 'logcamp_diag'
           d2 (str): The second data term; options are 'vis', 'bs', 'amp', 'cphase', 'cphase_diag' 'camp', 'logcamp', 'logcamp_diag'
           d3 (str): The third data term; options are 'vis', 'bs', 'amp', 'cphase', 'cphase_diag' 'camp', 'logcamp', 'logcamp_diag'

           s1 (str): The first regularizer; options are 'simple', 'gs', 'tv', 'tv2', 'l1', 'patch','compact','compact2','rgauss'
           s2 (str): The second regularizer; options are 'simple', 'gs', 'tv', 'tv2','l1', 'patch','compact','compact2','rgauss'
           s3 (str): The third regularizer; options are 'simple', 'gs', 'tv', 'tv2','l1', 'patch','compact','compact2','rgauss'

           alpha_d1 (float): The first data term weighting
           alpha_d2 (float): The second data term weighting
           alpha_d2 (float): The third data term weighting

           alpha_s1 (float): The first regularizer term weighting
           alpha_s2 (float): The second regularizer term weighting
           alpha_s3 (float): The third regularizer term weighting

           alpha_flux (float): The weighting for the total flux constraint
           alpha_cm (float): The weighting for the center of mass constraint

           maxit (int): Maximum number of minimizer iterations
           stop (float): The convergence criterion
           clipfloor (float): The Jy/pixel level above which prior image pixels are varied

           grads (bool): If True, analytic gradients are used
           logim (bool): If True, uses I = exp(I') change of variables
           norm_reg (bool): If True, normalizes regularizer terms
           norm_init (bool): If True, normalizes initial image to given total flux
           show_updates (bool): If True, displays the progress of the minimizer

           weighting (str): 'natural' or 'uniform'
           debias (bool): if True then apply debiasing to amplitudes/closure amplitudes
           systematic_noise (float): a fractional systematic noise tolerance to add to thermal sigmas
           snrcut (float): a  snr cutoff for including data in the chi^2 sum
           beam_size (float): beam size in radians for normalizing the regularizers

           maxset (bool):  if True, use maximal set instead of minimal for closure quantities
           systematic_cphase_noise (float): a value in degrees to add to the closure phase sigmas
           cp_uv_min (float): flag baselines shorter than this before forming closure quantities

           ttype (str): The Fourier transform type; options are 'fast', 'direct', 'nfft'
           fft_pad_factor (float): The FFT will pre-pad the image by this factor x the original size
           order (int): Interpolation order for sampling the FFT
           conv_func (str): The convolving function for gridding; options are 'gaussian', 'pill', and 'cubic'
           p_rad (int): The pixel radius for the convolving function in gridding for FFTs

       Returns:
           Image: Image object with result
    """

    # some kwarg default values
    maxit = kwargs.get('maxit', MAXIT)
    stop = kwargs.get('stop', STOP)
    clipfloor = kwargs.get('clipfloor', 0)
    ttype = kwargs.get('ttype', 'direct')

    grads = kwargs.get('grads', True)
    logim = kwargs.get('logim', True)
    norm_init = kwargs.get('norm_init', False)
    show_updates = kwargs.get('show_updates', True)

    beam_size = kwargs.get('beam_size', Obsdata.res())
    kwargs['beam_size'] = beam_size

    # Make sure data and regularizer options are ok
    if not d1 and not d2:
        raise Exception("Must have at least one data term!")
    if not s1 and not s2:
        raise Exception("Must have at least one regularizer term!")
    if (not ((d1 in DATATERMS) or d1 is False)) or (not ((d2 in DATATERMS) or d2 is False)):
        raise Exception("Invalid data term: valid data terms are: " + ' '.join(DATATERMS))
    if (not ((s1 in REGULARIZERS) or s1 is False)) or (not ((s2 in REGULARIZERS) or s2 is False)):
        raise Exception("Invalid regularizer: valid regularizers are: " + ' '.join(REGULARIZERS))
    if (Prior.psize != InitIm.psize) or (Prior.xdim != InitIm.xdim) or (Prior.ydim != InitIm.ydim):
        raise Exception("Initial image does not match dimensions of the prior image!")
    if (InitIm.polrep != Prior.polrep):
        raise Exception(
            "Initial image pol. representation does not match pol. representation of the prior image!")
    if (logim and Prior.pol_prim in ['Q', 'U', 'V']):
        raise Exception(
            "Cannot image Stokes Q,U,or V with log image transformation! Set logim=False in imager_func")

    pol = Prior.pol_prim
    print("Generating %s image..." % pol)

    # Catch scale and dimension problems
    imsize = np.max([Prior.xdim, Prior.ydim]) * Prior.psize
    uvmax = 1.0/Prior.psize
    uvmin = 1.0/imsize
    uvdists = Obsdata.unpack('uvdist')['uvdist']
    maxbl = np.max(uvdists)
    minbl = np.max(uvdists[uvdists > 0])
    maxamp = np.max(np.abs(Obsdata.unpack('amp')['amp']))

    if uvmax < maxbl:
        print("Warning! Pixel Spacing is larger than smallest spatial wavelength!")
    if uvmin > minbl:
        print("Warning! Field of View is smaller than largest nonzero spatial wavelength!")
    if flux > 1.2*maxamp:
        print("Warning! Specified flux is > 120% of maximum visibility amplitude!")
    if flux < .8*maxamp:
        print("Warning! Specified flux is < 80% of maximum visibility amplitude!")

    # Define embedding mask
    embed_mask = Prior.imvec > clipfloor

    # Normalize prior image to total flux and limit imager range to prior values > clipfloor
    if (not norm_init):
        nprior = Prior.imvec[embed_mask]
        ninit = InitIm.imvec[embed_mask]
    else:
        nprior = (flux * Prior.imvec / np.sum((Prior.imvec)[embed_mask]))[embed_mask]
        ninit = (flux * InitIm.imvec / np.sum((InitIm.imvec)[embed_mask]))[embed_mask]

    if len(nprior) == 0:
        raise Exception("clipfloor too large: all prior pixels have been clipped!")

    # Get data and fourier matrices for the data terms
    (data1, sigma1, A1) = chisqdata(Obsdata, Prior, embed_mask, d1, pol=pol, **kwargs)
    (data2, sigma2, A2) = chisqdata(Obsdata, Prior, embed_mask, d2, pol=pol, **kwargs)
    (data3, sigma3, A3) = chisqdata(Obsdata, Prior, embed_mask, d3, pol=pol, **kwargs)

    # Define the chi^2 and chi^2 gradient
    def chisq1(imvec):
        return chisq(imvec, A1, data1, sigma1, d1, ttype=ttype, mask=embed_mask)

    def chisq1grad(imvec):
        c = chisqgrad(imvec, A1, data1, sigma1, d1, ttype=ttype, mask=embed_mask)
        return c

    def chisq2(imvec):
        return chisq(imvec, A2, data2, sigma2, d2, ttype=ttype, mask=embed_mask)

    def chisq2grad(imvec):
        c = chisqgrad(imvec, A2, data2, sigma2, d2, ttype=ttype, mask=embed_mask)
        return c

    def chisq3(imvec):
        return chisq(imvec, A3, data3, sigma3, d3, ttype=ttype, mask=embed_mask)

    def chisq3grad(imvec):
        c = chisqgrad(imvec, A3, data3, sigma3, d3, ttype=ttype, mask=embed_mask)
        return c

    # Define the regularizer and regularizer gradient
    def reg1(imvec):
        return regularizer(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, s1, **kwargs)

    def reg1grad(imvec):
        return regularizergrad(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, s1, **kwargs)

    def reg2(imvec):
        return regularizer(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, s2, **kwargs)

    def reg2grad(imvec):
        return regularizergrad(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, s2, **kwargs)

    def reg3(imvec):
        return regularizer(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, s3, **kwargs)

    def reg3grad(imvec):
        return regularizergrad(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, s3, **kwargs)

    # Define constraint functions
    def flux_constraint(imvec):
        return regularizer(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, "flux", **kwargs)

    def flux_constraint_grad(imvec):
        return regularizergrad(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, "flux", **kwargs)

    def cm_constraint(imvec):
        return regularizer(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, "cm", **kwargs)

    def cm_constraint_grad(imvec):
        return regularizergrad(imvec, nprior, embed_mask, flux, Prior.xdim, Prior.ydim, Prior.psize, "cm", **kwargs)

    # Define the objective function and gradient
    def objfunc(imvec):
        if logim:
            imvec = np.exp(imvec)

        datterm = alpha_d1 * (chisq1(imvec) - 1) + alpha_d2 * \
            (chisq2(imvec) - 1) + alpha_d3 * (chisq3(imvec) - 1)
        regterm = alpha_s1 * reg1(imvec) + alpha_s2 * reg2(imvec) + alpha_s3 * reg3(imvec)
        conterm = alpha_flux * flux_constraint(imvec) + alpha_cm * cm_constraint(imvec)

        return datterm + regterm + conterm

    def objgrad(imvec):
        if logim:
            imvec = np.exp(imvec)

        datterm = alpha_d1 * chisq1grad(imvec) + alpha_d2 * \
            chisq2grad(imvec) + alpha_d3 * chisq3grad(imvec)
        regterm = alpha_s1 * reg1grad(imvec) + alpha_s2 * \
            reg2grad(imvec) + alpha_s3 * reg3grad(imvec)
        conterm = alpha_flux * flux_constraint_grad(imvec) + alpha_cm * cm_constraint_grad(imvec)

        grad = datterm + regterm + conterm

        # chain rule term for change of variables
        if logim:
            grad *= imvec

        return grad

    # Define plotting function for each iteration
    global nit
    nit = 0

    def plotcur(im_step):
        global nit
        if logim:
            im_step = np.exp(im_step)
        if show_updates:
            chi2_1 = chisq1(im_step)
            chi2_2 = chisq2(im_step)
            chi2_3 = chisq3(im_step)
            s_1 = reg1(im_step)
            s_2 = reg2(im_step)
            s_3 = reg3(im_step)
            if np.any(np.invert(embed_mask)):
                im_step = embed(im_step, embed_mask)
            plot_i(im_step, Prior, nit, {d1: chi2_1, d2: chi2_2, d3: chi2_3}, pol=pol)
            print("i: %d chi2_1: %0.2f chi2_2: %0.2f chi2_3: %0.2f s_1: %0.2f s_2: %0.2f s_3: %0.2f" % (
                nit, chi2_1, chi2_2, chi2_3, s_1, s_2, s_3))
        nit += 1

    # Generate and the initial image
    if logim:
        xinit = np.log(ninit)
    else:
        xinit = ninit

    # Print stats
    print("Initial S_1: %f S_2: %f S_3: %f" % (reg1(ninit), reg2(ninit), reg3(ninit)))
    print("Initial Chi^2_1: %f Chi^2_2: %f Chi^2_3: %f" %
          (chisq1(ninit), chisq2(ninit), chisq3(ninit)))
    print("Initial Objective Function: %f" % (objfunc(xinit)))

    if d1 in DATATERMS:
        print("Total Data 1: ", (len(data1)))
    if d2 in DATATERMS:
        print("Total Data 2: ", (len(data2)))
    if d3 in DATATERMS:
        print("Total Data 3: ", (len(data3)))

    print("Total Pixel #: ", (len(Prior.imvec)))
    print("Clipped Pixel #: ", (len(ninit)))
    print()
    plotcur(xinit)

    # Minimize
    optdict = {'maxiter': maxit, 'ftol': stop, 'maxcor': NHIST,
               'gtol': stop, 'maxls': MAXLS}  # minimizer dict params
    tstart = time.time()
    if grads:
        res = opt.minimize(objfunc, xinit, method='L-BFGS-B', jac=objgrad,
                           options=optdict, callback=plotcur)
    else:
        res = opt.minimize(objfunc, xinit, method='L-BFGS-B',
                           options=optdict, callback=plotcur)

    tstop = time.time()

    # Format output
    out = res.x
    if logim:
        out = np.exp(res.x)
    if np.any(np.invert(embed_mask)):
        out = embed(out, embed_mask)

    outim = image.Image(out.reshape(Prior.ydim, Prior.xdim),
                        Prior.psize, Prior.ra, Prior.dec,
                        rf=Prior.rf, source=Prior.source,
                        polrep=Prior.polrep, pol_prim=pol,
                        mjd=Prior.mjd, time=Prior.time, pulse=Prior.pulse)

    # copy over other polarizations
    outim.copy_pol_images(InitIm)

    # Print stats
    print("time: %f s" % (tstop - tstart))
    print("J: %f" % res.fun)
    print("Final Chi^2_1: %f Chi^2_2: %f  Chi^2_3: %f" %
          (chisq1(out[embed_mask]), chisq2(out[embed_mask]), chisq3(out[embed_mask])))
    print(res.message)

    # Return Image object
    return outim

##################################################################################################
# Wrapper Functions
##################################################################################################


def chisq(imvec, A, data, sigma, dtype, ttype='direct', mask=None):
    """return the chi^2 for the appropriate dtype
    """

    if mask is None:
        mask = []
    chisq = 1
    if dtype not in DATATERMS:
        return chisq

    if ttype not in ['fast', 'direct', 'nfft']:
        raise Exception("Possible ttype values are 'fast', 'direct'!, 'nfft!'")

    if ttype == 'direct':
        if dtype == 'vis':
            chisq = chisq_vis(imvec, A, data, sigma)
        elif dtype == 'amp':
            chisq = chisq_amp(imvec, A, data, sigma)
        elif dtype == 'logamp':
            chisq = chisq_logamp(imvec, A, data, sigma)
        elif dtype == 'bs':
            chisq = chisq_bs(imvec, A, data, sigma)
        elif dtype == 'cphase':
            chisq = chisq_cphase(imvec, A, data, sigma)
        elif dtype == 'cphase_diag':
            chisq = chisq_cphase_diag(imvec, A, data, sigma)
        elif dtype == 'camp':
            chisq = chisq_camp(imvec, A, data, sigma)
        elif dtype == 'logcamp':
            chisq = chisq_logcamp(imvec, A, data, sigma)
        elif dtype == 'logcamp_diag':
            chisq = chisq_logcamp_diag(imvec, A, data, sigma)

    elif ttype == 'fast':
        if len(mask) > 0 and np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)

        if dtype not in ['cphase_diag', 'logcamp_diag']:
            vis_arr = obsh.fft_imvec(imvec, A[0])

        if dtype == 'vis':
            chisq = chisq_vis_fft(vis_arr, A, data, sigma)
        elif dtype == 'amp':
            chisq = chisq_amp_fft(vis_arr, A, data, sigma)
        elif dtype == 'logamp':
            chisq = chisq_logamp_fft(vis_arr, A, data, sigma)
        elif dtype == 'bs':
            chisq = chisq_bs_fft(vis_arr, A, data, sigma)
        elif dtype == 'cphase':
            chisq = chisq_cphase_fft(vis_arr, A, data, sigma)
        elif dtype == 'cphase_diag':
            chisq = chisq_cphase_diag_fft(imvec, A, data, sigma)
        elif dtype == 'camp':
            chisq = chisq_camp_fft(vis_arr, A, data, sigma)
        elif dtype == 'logcamp':
            chisq = chisq_logcamp_fft(vis_arr, A, data, sigma)
        elif dtype == 'logcamp_diag':
            chisq = chisq_logcamp_diag_fft(imvec, A, data, sigma)

    elif ttype == 'nfft':
        if len(mask) > 0 and np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)

        if dtype == 'vis':
            chisq = chisq_vis_nfft(imvec, A, data, sigma)
        elif dtype == 'amp':
            chisq = chisq_amp_nfft(imvec, A, data, sigma)
        elif dtype == 'logamp':
            chisq = chisq_logamp_nfft(imvec, A, data, sigma)
        elif dtype == 'bs':
            chisq = chisq_bs_nfft(imvec, A, data, sigma)
        elif dtype == 'cphase':
            chisq = chisq_cphase_nfft(imvec, A, data, sigma)
        elif dtype == 'cphase_diag':
            chisq = chisq_cphase_diag_nfft(imvec, A, data, sigma)
        elif dtype == 'camp':
            chisq = chisq_camp_nfft(imvec, A, data, sigma)
        elif dtype == 'logcamp':
            chisq = chisq_logcamp_nfft(imvec, A, data, sigma)
        elif dtype == 'logcamp_diag':
            chisq = chisq_logcamp_diag_nfft(imvec, A, data, sigma)

    return chisq


def chisqgrad(imvec, A, data, sigma, dtype, ttype='direct', mask=None):
    """return the chi^2 gradient for the appropriate dtype
    """

    if mask is None:
        mask = []
    chisqgrad = np.zeros(len(imvec))
    if dtype not in DATATERMS:
        return chisqgrad

    if ttype not in ['fast', 'direct', 'nfft']:
        raise Exception("Possible ttype values are 'fast', 'direct', 'nfft'!")

    if ttype == 'direct':
        if dtype == 'vis':
            chisqgrad = chisqgrad_vis(imvec, A, data, sigma)
        elif dtype == 'amp':
            chisqgrad = chisqgrad_amp(imvec, A, data, sigma)
        elif dtype == 'logamp':
            chisqgrad = chisqgrad_logamp(imvec, A, data, sigma)
        elif dtype == 'bs':
            chisqgrad = chisqgrad_bs(imvec, A, data, sigma)
        elif dtype == 'cphase':
            chisqgrad = chisqgrad_cphase(imvec, A, data, sigma)
        elif dtype == 'cphase_diag':
            chisqgrad = chisqgrad_cphase_diag(imvec, A, data, sigma)
        elif dtype == 'camp':
            chisqgrad = chisqgrad_camp(imvec, A, data, sigma)
        elif dtype == 'logcamp':
            chisqgrad = chisqgrad_logcamp(imvec, A, data, sigma)
        elif dtype == 'logcamp_diag':
            chisqgrad = chisqgrad_logcamp_diag(imvec, A, data, sigma)

    elif ttype == 'fast':
        if len(mask) > 0 and np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)

        if dtype not in ['cphase_diag', 'logcamp_diag']:
            vis_arr = obsh.fft_imvec(imvec, A[0])

        if dtype == 'vis':
            chisqgrad = chisqgrad_vis_fft(vis_arr, A, data, sigma)
        elif dtype == 'amp':
            chisqgrad = chisqgrad_amp_fft(vis_arr, A, data, sigma)
        elif dtype == 'logamp':
            chisqgrad = chisqgrad_logamp_fft(vis_arr, A, data, sigma)
        elif dtype == 'bs':
            chisqgrad = chisqgrad_bs_fft(vis_arr, A, data, sigma)
        elif dtype == 'cphase':
            chisqgrad = chisqgrad_cphase_fft(vis_arr, A, data, sigma)
        elif dtype == 'cphase_diag':
            chisqgrad = chisqgrad_cphase_diag_fft(imvec, A, data, sigma)
        elif dtype == 'camp':
            chisqgrad = chisqgrad_camp_fft(vis_arr, A, data, sigma)
        elif dtype == 'logcamp':
            chisqgrad = chisqgrad_logcamp_fft(vis_arr, A, data, sigma)
        elif dtype == 'logcamp_diag':
            chisqgrad = chisqgrad_logcamp_diag_fft(imvec, A, data, sigma)

        if len(mask) > 0 and np.any(np.invert(mask)):
            chisqgrad = chisqgrad[mask]

    elif ttype == 'nfft':
        if len(mask) > 0 and np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)

        if dtype == 'vis':
            chisqgrad = chisqgrad_vis_nfft(imvec, A, data, sigma)
        elif dtype == 'amp':
            chisqgrad = chisqgrad_amp_nfft(imvec, A, data, sigma)
        elif dtype == 'logamp':
            chisqgrad = chisqgrad_logamp_nfft(imvec, A, data, sigma)
        elif dtype == 'bs':
            chisqgrad = chisqgrad_bs_nfft(imvec, A, data, sigma)
        elif dtype == 'cphase':
            chisqgrad = chisqgrad_cphase_nfft(imvec, A, data, sigma)
        elif dtype == 'cphase_diag':
            chisqgrad = chisqgrad_cphase_diag_nfft(imvec, A, data, sigma)
        elif dtype == 'camp':
            chisqgrad = chisqgrad_camp_nfft(imvec, A, data, sigma)
        elif dtype == 'logcamp':
            chisqgrad = chisqgrad_logcamp_nfft(imvec, A, data, sigma)
        elif dtype == 'logcamp_diag':
            chisqgrad = chisqgrad_logcamp_diag_nfft(imvec, A, data, sigma)

        if len(mask) > 0 and np.any(np.invert(mask)):
            chisqgrad = chisqgrad[mask]

    return chisqgrad


def regularizer(imvec, nprior, mask, flux, xdim, ydim, psize, stype, **kwargs):
    """return the regularizer value
    """

    norm_reg = kwargs.get('norm_reg', NORM_REGULARIZER)
    beam_size = kwargs.get('beam_size', psize)
    alpha_A = kwargs.get('alpha_A', 1.0)
    epsilon = kwargs.get('epsilon_tv', 0.)

    if stype == "flux":
        s = -sflux(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "cm":
        s = -scm(imvec, xdim, ydim, psize, flux, mask, norm_reg=norm_reg, beam_size=beam_size)
    elif stype == "simple":
        s = -ssimple(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "l1":
        s = -sl1(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "l1w":
        s = -sl1w(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "lA":
        s = -slA(imvec, nprior, psize, flux, beam_size, alpha_A, norm_reg)
    elif stype == "gs":
        s = -sgs(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "patch":
        s = -spatch(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "tv":
        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -stv(imvec, xdim, ydim, psize, flux, norm_reg=norm_reg,
                 beam_size=beam_size, epsilon=epsilon)
    elif stype == "tv2":
        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -stv2(imvec, xdim, ydim, psize, flux, norm_reg=norm_reg, beam_size=beam_size)
    elif stype == "compact":
        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -scompact(imvec, xdim, ydim, psize, flux, norm_reg=norm_reg)
    elif stype == "compact2":
        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -scompact2(imvec, xdim, ydim, psize, flux, norm_reg=norm_reg)
    elif stype == "rgauss":
        # additional key words for gaussian regularizer
        major = kwargs.get('major', 1.0)
        minor = kwargs.get('minor', 1.0)
        PA = kwargs.get('PA', 1.0)

        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -sgauss(imvec, xdim, ydim, psize, major=major, minor=minor, PA=PA)
    else:
        s = 0

    return s


def regularizergrad(imvec, nprior, mask, flux, xdim, ydim, psize, stype, **kwargs):
    """return the regularizer gradient
    """

    norm_reg = kwargs.get('norm_reg', NORM_REGULARIZER)
    beam_size = kwargs.get('beam_size', psize)
    alpha_A = kwargs.get('alpha_A', 1.0)
    epsilon = kwargs.get('epsilon_tv', 0.)

    if stype == "flux":
        s = -sfluxgrad(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "cm":
        s = -scmgrad(imvec, xdim, ydim, psize, flux, mask, norm_reg=norm_reg, beam_size=beam_size)
    elif stype == "simple":
        s = -ssimplegrad(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "l1":
        s = -sl1grad(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "l1w":
        s = -sl1wgrad(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "lA":
        s = -slAgrad(imvec, nprior, psize, flux, beam_size, alpha_A, norm_reg)
    elif stype == "gs":
        s = -sgsgrad(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "patch":
        s = -spatchgrad(imvec, nprior, flux, norm_reg=norm_reg)
    elif stype == "tv":
        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -stvgrad(imvec, xdim, ydim, psize, flux, norm_reg=norm_reg,
                     beam_size=beam_size, epsilon=epsilon)
        s = s[mask]
    elif stype == "tv2":
        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -stv2grad(imvec, xdim, ydim, psize, flux, norm_reg=norm_reg, beam_size=beam_size)
        s = s[mask]
    elif stype == "compact":
        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -scompactgrad(imvec, xdim, ydim, psize, flux, norm_reg=norm_reg)
        s = s[mask]
    elif stype == "compact2":
        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -scompact2grad(imvec, xdim, ydim, psize, flux, norm_reg=norm_reg)
        s = s[mask]
    elif stype == "rgauss":
        # additional key words for gaussian regularizer
        major = kwargs.get('major', 1.0)
        minor = kwargs.get('minor', 1.0)
        PA = kwargs.get('PA', 1.0)

        if np.any(np.invert(mask)):
            imvec = embed(imvec, mask, randomfloor=True)
        s = -sgauss_grad(imvec, xdim, ydim, psize, major, minor, PA)
        s = s[mask]
    else:
        s = np.zeros(len(imvec))

    return s


def chisqdata(Obsdata, Prior, mask, dtype, pol='I', **kwargs):
    """Return the data, sigma, and matrices for the appropriate dtype
    """

    ttype = kwargs.get('ttype', 'direct')
    (data, sigma, A) = (False, False, False)
    if ttype not in ['fast', 'direct', 'nfft']:
        raise Exception("Possible ttype values are 'fast', 'direct', 'nfft'!")

    if ttype == 'direct':
        if dtype == 'vis':
            (data, sigma, A) = chisqdata_vis(Obsdata, Prior, mask, pol=pol, **kwargs)
        elif dtype == 'amp' or dtype == 'logamp':
            (data, sigma, A) = chisqdata_amp(Obsdata, Prior, mask, pol=pol, **kwargs)
        elif dtype == 'bs':
            (data, sigma, A) = chisqdata_bs(Obsdata, Prior, mask, pol=pol, **kwargs)
        elif dtype == 'cphase':
            (data, sigma, A) = chisqdata_cphase(Obsdata, Prior, mask, pol=pol, **kwargs)
        elif dtype == 'cphase_diag':
            (data, sigma, A) = chisqdata_cphase_diag(Obsdata, Prior, mask, pol=pol, **kwargs)
        elif dtype == 'camp':
            (data, sigma, A) = chisqdata_camp(Obsdata, Prior, mask, pol=pol, **kwargs)
        elif dtype == 'logcamp':
            (data, sigma, A) = chisqdata_logcamp(Obsdata, Prior, mask, pol=pol, **kwargs)
        elif dtype == 'logcamp_diag':
            (data, sigma, A) = chisqdata_logcamp_diag(Obsdata, Prior, mask, pol=pol, **kwargs)

    elif ttype == 'fast':
        if dtype == 'vis':
            (data, sigma, A) = chisqdata_vis_fft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'amp' or dtype == 'logamp':
            (data, sigma, A) = chisqdata_amp_fft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'bs':
            (data, sigma, A) = chisqdata_bs_fft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'cphase':
            (data, sigma, A) = chisqdata_cphase_fft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'cphase_diag':
            (data, sigma, A) = chisqdata_cphase_diag_fft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'camp':
            (data, sigma, A) = chisqdata_camp_fft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'logcamp':
            (data, sigma, A) = chisqdata_logcamp_fft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'logcamp_diag':
            (data, sigma, A) = chisqdata_logcamp_diag_fft(Obsdata, Prior, pol=pol, **kwargs)

    elif ttype == 'nfft':
        if dtype == 'vis':
            (data, sigma, A) = chisqdata_vis_nfft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'amp' or dtype == 'logamp':
            (data, sigma, A) = chisqdata_amp_nfft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'bs':
            (data, sigma, A) = chisqdata_bs_nfft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'cphase':
            (data, sigma, A) = chisqdata_cphase_nfft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'cphase_diag':
            (data, sigma, A) = chisqdata_cphase_diag_nfft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'camp':
            (data, sigma, A) = chisqdata_camp_nfft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'logcamp':
            (data, sigma, A) = chisqdata_logcamp_nfft(Obsdata, Prior, pol=pol, **kwargs)
        elif dtype == 'logcamp_diag':
            (data, sigma, A) = chisqdata_logcamp_diag_nfft(Obsdata, Prior, pol=pol, **kwargs)

    return (data, sigma, A)


##################################################################################################
# DFT Chi-squared and Gradient Functions
##################################################################################################

def chisq_vis(imvec, Amatrix, vis, sigma):
    """Visibility chi-squared"""

    samples = np.dot(Amatrix, imvec)
    chisq = np.sum(np.abs((samples-vis)/sigma)**2)/(2*len(vis))
    return chisq

def chisqgrad_vis(imvec, Amatrix, vis, sigma):
    """The gradient of the visibility chi-squared"""

    samples = np.dot(Amatrix, imvec)
    wdiff = (vis - samples)/(sigma**2)

    out = -np.real(np.dot(Amatrix.conj().T, wdiff))/len(vis)
    return out


def chisq_amp(imvec, A, amp, sigma):
    """Visibility Amplitudes (normalized) chi-squared"""

    amp_samples = np.abs(np.dot(A, imvec))
    return np.sum(np.abs((amp - amp_samples)/sigma)**2)/len(amp)


def chisqgrad_amp(imvec, A, amp, sigma):
    """The gradient of the amplitude chi-squared"""

    i1 = np.dot(A, imvec)
    amp_samples = np.abs(i1)

    pp = ((amp - amp_samples) * amp_samples) / (sigma**2) / i1
    out = (-2.0/len(amp)) * np.real(np.dot(pp, A))
    return out


def chisq_bs(imvec, Amatrices, bis, sigma):
    """Bispectrum chi-squared"""

    bisamples = (np.dot(Amatrices[0], imvec) *
                 np.dot(Amatrices[1], imvec) *
                 np.dot(Amatrices[2], imvec))
    chisq = np.sum(np.abs(((bis - bisamples)/sigma))**2)/(2.*len(bis))
    return chisq


def chisqgrad_bs(imvec, Amatrices, bis, sigma):
    """The gradient of the bispectrum chi-squared"""

    bisamples = (np.dot(Amatrices[0], imvec) *
                 np.dot(Amatrices[1], imvec) *
                 np.dot(Amatrices[2], imvec))

    wdiff = ((bis - bisamples).conj())/(sigma**2)
    pt1 = wdiff * np.dot(Amatrices[1], imvec) * np.dot(Amatrices[2], imvec)
    pt2 = wdiff * np.dot(Amatrices[0], imvec) * np.dot(Amatrices[2], imvec)
    pt3 = wdiff * np.dot(Amatrices[0], imvec) * np.dot(Amatrices[1], imvec)
    out = (np.dot(pt1, Amatrices[0]) +
           np.dot(pt2, Amatrices[1]) + 
           np.dot(pt3, Amatrices[2]))

    out = -np.real(out) / len(bis)
    return out


def chisq_cphase(imvec, Amatrices, clphase, sigma):
    """Closure Phases (normalized) chi-squared"""
    clphase = clphase * ehc.DEGREE
    sigma = sigma * ehc.DEGREE

    i1 = np.dot(Amatrices[0], imvec)
    i2 = np.dot(Amatrices[1], imvec)
    i3 = np.dot(Amatrices[2], imvec)
    clphase_samples = np.angle(i1 * i2 * i3)

    chisq = (2.0/len(clphase)) * np.sum((1.0 - np.cos(clphase-clphase_samples))/(sigma**2))
    return chisq


def chisqgrad_cphase(imvec, Amatrices, clphase, sigma):
    """The gradient of the closure phase chi-squared"""
    clphase = clphase * ehc.DEGREE
    sigma = sigma * ehc.DEGREE

    i1 = np.dot(Amatrices[0], imvec)
    i2 = np.dot(Amatrices[1], imvec)
    i3 = np.dot(Amatrices[2], imvec)
    clphase_samples = np.angle(i1 * i2 * i3)

    pref = np.sin(clphase - clphase_samples)/(sigma**2)
    pt1 = pref/i1
    pt2 = pref/i2
    pt3 = pref/i3
    out = np.dot(pt1, Amatrices[0]) + np.dot(pt2, Amatrices[1]) + np.dot(pt3, Amatrices[2])
    out = (-2.0/len(clphase)) * np.imag(out)
    return out


def chisq_cphase_diag(imvec, Amatrices, clphase_diag, sigma):
    """Diagonalized closure phases (normalized) chi-squared"""
    clphase_diag = np.concatenate(clphase_diag) * ehc.DEGREE
    sigma = np.concatenate(sigma) * ehc.DEGREE

    A3_diag = Amatrices[0]
    tform_mats = Amatrices[1]

    clphase_diag_samples = []
    for iA, A3 in enumerate(A3_diag):
        clphase_samples = np.angle(np.dot(A3[0], imvec) *
                                   np.dot(A3[1], imvec) *
                                   np.dot(A3[2], imvec))
        clphase_diag_samples.append(np.dot(tform_mats[iA], clphase_samples))
    clphase_diag_samples = np.concatenate(clphase_diag_samples)

    chisq = np.sum((1.0 - np.cos(clphase_diag-clphase_diag_samples))/(sigma**2))
    chisq *= (2.0/len(clphase_diag))
    return chisq


def chisqgrad_cphase_diag(imvec, Amatrices, clphase_diag, sigma):
    """The gradient of the diagonalized closure phase chi-squared"""
    clphase_diag = clphase_diag * ehc.DEGREE
    sigma = sigma * ehc.DEGREE

    A3_diag = Amatrices[0]
    tform_mats = Amatrices[1]

    deriv = np.zeros_like(imvec)
    for iA, A3 in enumerate(A3_diag):

        i1 = np.dot(A3[0], imvec)
        i2 = np.dot(A3[1], imvec)
        i3 = np.dot(A3[2], imvec)
        clphase_samples = np.angle(i1 * i2 * i3)
        clphase_diag_samples = np.dot(tform_mats[iA], clphase_samples)

        clphase_diag_measured = clphase_diag[iA]
        clphase_diag_sigma = sigma[iA]

        term1 = np.dot(np.dot((np.sin(clphase_diag_measured-clphase_diag_samples) /
                               (clphase_diag_sigma**2.0)), (tform_mats[iA]/i1)), A3[0])
        term2 = np.dot(np.dot((np.sin(clphase_diag_measured-clphase_diag_samples) /
                               (clphase_diag_sigma**2.0)), (tform_mats[iA]/i2)), A3[1])
        term3 = np.dot(np.dot((np.sin(clphase_diag_measured-clphase_diag_samples) /
                               (clphase_diag_sigma**2.0)), (tform_mats[iA]/i3)), A3[2])
        deriv += -2.0*np.imag(term1 + term2 + term3)

    deriv *= 1.0/np.float(len(np.concatenate(clphase_diag)))

    return deriv


def chisq_camp(imvec, Amatrices, clamp, sigma):
    """Closure Amplitudes (normalized) chi-squared"""

    i1 = np.dot(Amatrices[0], imvec)
    i2 = np.dot(Amatrices[1], imvec)
    i3 = np.dot(Amatrices[2], imvec)
    i4 = np.dot(Amatrices[3], imvec)
    clamp_samples = np.abs((i1 * i2)/(i3 * i4))

    chisq = np.sum(np.abs((clamp - clamp_samples)/sigma)**2)/len(clamp)
    return chisq


def chisqgrad_camp(imvec, Amatrices, clamp, sigma):
    """The gradient of the closure amplitude chi-squared"""

    i1 = np.dot(Amatrices[0], imvec)
    i2 = np.dot(Amatrices[1], imvec)
    i3 = np.dot(Amatrices[2], imvec)
    i4 = np.dot(Amatrices[3], imvec)
    clamp_samples = np.abs((i1 * i2)/(i3 * i4))

    pp = ((clamp - clamp_samples) * clamp_samples)/(sigma**2)
    pt1 = pp/i1
    pt2 = pp/i2
    pt3 = -pp/i3
    pt4 = -pp/i4
    out = (np.dot(pt1, Amatrices[0]) +
           np.dot(pt2, Amatrices[1]) +
           np.dot(pt3, Amatrices[2]) +
           np.dot(pt4, Amatrices[3]))
    out *= (-2.0/len(clamp)) * np.real(out)
    return out


def chisq_logcamp(imvec, Amatrices, log_clamp, sigma):
    """Log Closure Amplitudes (normalized) chi-squared"""

    a1 = np.abs(np.dot(Amatrices[0], imvec))
    a2 = np.abs(np.dot(Amatrices[1], imvec))
    a3 = np.abs(np.dot(Amatrices[2], imvec))
    a4 = np.abs(np.dot(Amatrices[3], imvec))

    samples = np.log(a1) + np.log(a2) - np.log(a3) - np.log(a4)
    chisq = np.sum(np.abs((log_clamp - samples)/sigma)**2) / (len(log_clamp))
    return chisq


def chisqgrad_logcamp(imvec, Amatrices, log_clamp, sigma):
    """The gradient of the Log closure amplitude chi-squared"""

    i1 = np.dot(Amatrices[0], imvec)
    i2 = np.dot(Amatrices[1], imvec)
    i3 = np.dot(Amatrices[2], imvec)
    i4 = np.dot(Amatrices[3], imvec)
    log_clamp_samples = (np.log(np.abs(i1)) +
                         np.log(np.abs(i2)) - 
                         np.log(np.abs(i3)) -
                         np.log(np.abs(i4)))

    pp = (log_clamp - log_clamp_samples) / (sigma**2)
    pt1 = pp / i1
    pt2 = pp / i2
    pt3 = -pp / i3
    pt4 = -pp / i4
    out = (np.dot(pt1, Amatrices[0]) +
           np.dot(pt2, Amatrices[1]) +
           np.dot(pt3, Amatrices[2]) +
           np.dot(pt4, Amatrices[3]))
    out = (-2.0/len(log_clamp)) * np.real(out)
    return out


def chisq_logcamp_diag(imvec, Amatrices, log_clamp_diag, sigma):
    """Diagonalized log closure amplitudes (normalized) chi-squared"""

    log_clamp_diag = np.concatenate(log_clamp_diag)
    sigma = np.concatenate(sigma)

    A4_diag = Amatrices[0]
    tform_mats = Amatrices[1]

    log_clamp_diag_samples = []
    for iA, A4 in enumerate(A4_diag):

        a1 = np.abs(np.dot(A4[0], imvec))
        a2 = np.abs(np.dot(A4[1], imvec))
        a3 = np.abs(np.dot(A4[2], imvec))
        a4 = np.abs(np.dot(A4[3], imvec))

        log_clamp_samples = np.log(a1) + np.log(a2) - np.log(a3) - np.log(a4)
        log_clamp_diag_samples.append(np.dot(tform_mats[iA], log_clamp_samples))

    log_clamp_diag_samples = np.concatenate(log_clamp_diag_samples)

    chisq = np.sum(np.abs((log_clamp_diag - log_clamp_diag_samples)/sigma)**2)
    chisq /= (len(log_clamp_diag))

    return chisq


def chisqgrad_logcamp_diag(imvec, Amatrices, log_clamp_diag, sigma):
    """The gradient of the diagonalized log closure amplitude chi-squared"""

    A4_diag = Amatrices[0]
    tform_mats = Amatrices[1]

    deriv = np.zeros_like(imvec)
    for iA, A4 in enumerate(A4_diag):

        i1 = np.dot(A4[0], imvec)
        i2 = np.dot(A4[1], imvec)
        i3 = np.dot(A4[2], imvec)
        i4 = np.dot(A4[3], imvec)
        log_clamp_samples = np.log(np.abs(i1)) + np.log(np.abs(i2)) - \
            np.log(np.abs(i3)) - np.log(np.abs(i4))
        log_clamp_diag_samples = np.dot(tform_mats[iA], log_clamp_samples)

        log_clamp_diag_measured = log_clamp_diag[iA]
        log_clamp_diag_sigma = sigma[iA]

        term1 = np.dot(np.dot(((log_clamp_diag_measured-log_clamp_diag_samples) /
                               (log_clamp_diag_sigma**2.0)), (tform_mats[iA]/i1)), A4[0])
        term2 = np.dot(np.dot(((log_clamp_diag_measured-log_clamp_diag_samples) /
                               (log_clamp_diag_sigma**2.0)), (tform_mats[iA]/i2)), A4[1])
        term3 = np.dot(np.dot(((log_clamp_diag_measured-log_clamp_diag_samples) /
                               (log_clamp_diag_sigma**2.0)), (tform_mats[iA]/i3)), A4[2])
        term4 = np.dot(np.dot(((log_clamp_diag_measured-log_clamp_diag_samples) /
                               (log_clamp_diag_sigma**2.0)), (tform_mats[iA]/i4)), A4[3])
        deriv += -2.0*np.real(term1 + term2 - term3 - term4)

    deriv *= 1.0/np.float(len(np.concatenate(log_clamp_diag)))

    return deriv


def chisq_logamp(imvec, A, amp, sigma):
    """Log Visibility Amplitudes (normalized) chi-squared"""

    # to lowest order the variance on the logarithm of a quantity x is
    # sigma^2_log(x) = sigma^2/x^2
    logsigma = sigma / amp

    amp_samples = np.abs(np.dot(A, imvec))
    chisq = np.sum(np.abs((np.log(amp) - np.log(amp_samples))/logsigma)**2)/len(amp)
    return chisq

def chisqgrad_logamp(imvec, A, amp, sigma):
    """The gradient of the Log amplitude chi-squared"""

    # to lowest order the variance on the logarithm of a quantity x is
    # sigma^2_log(x) = sigma^2/x^2
    logsigma = sigma / amp

    i1 = np.dot(A, imvec)
    amp_samples = np.abs(i1)

    pp = ((np.log(amp) - np.log(amp_samples))) / (logsigma**2) / i1
    out = (-2.0/len(amp)) * np.real(np.dot(pp, A))
    return out

##################################################################################################
# FFT Chi-squared and Gradient Functions
##################################################################################################


def chisq_vis_fft(vis_arr, A, vis, sigma):
    """Visibility chi-squared from fft
    """

    im_info, sampler_info_list, gridder_info_list = A
    samples = obsh.sampler(vis_arr, sampler_info_list, sample_type="vis")

    chisq = np.sum(np.abs((samples-vis)/sigma)**2)/(2*len(vis))

    return chisq


def chisqgrad_vis_fft(vis_arr, A, vis, sigma):
    """The gradient of the visibility chi-squared from fft
    """

    im_info, sampler_info_list, gridder_info_list = A

    # samples and gradient FT
    pulsefac = sampler_info_list[0].pulsefac
    samples = obsh.sampler(vis_arr, sampler_info_list, sample_type="vis")
    wdiff_vec = (-1.0/len(vis)*(vis - samples)/(sigma**2)) * pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff_arr = obsh.gridder([wdiff_vec], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff_arr)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevant cells and flatten
    # TODO or is x<-->y??
    out = np.real(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                           im_info.padvaly1:-im_info.padvaly2].flatten())

    return out


def chisq_amp_fft(vis_arr, A, amp, sigma):
    """Visibility amplitude chi-squared from fft
    """

    im_info, sampler_info_list, gridder_info_list = A
    amp_samples = np.abs(obsh.sampler(vis_arr, sampler_info_list, sample_type="vis"))
    chisq = np.sum(np.abs((amp_samples-amp)/sigma)**2)/(len(amp))
    return chisq


def chisqgrad_amp_fft(vis_arr, A, amp, sigma):
    """The gradient of the amplitude chi-kernesquared
    """

    im_info, sampler_info_list, gridder_info_list = A

    # samples
    samples = obsh.sampler(vis_arr, sampler_info_list, sample_type="vis")
    amp_samples = np.abs(samples)

    # gradient FT
    pulsefac = sampler_info_list[0].pulsefac
    wdiff_vec = (-2.0/len(amp)*((amp - amp_samples) * amp_samples) /
                 (sigma**2) / samples.conj()) * pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff_arr = obsh.gridder([wdiff_vec], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff_arr)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevent cells and flatten
    # TODO or is x<-->y??
    out = np.real(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                           im_info.padvaly1:-im_info.padvaly2].flatten())

    return out


def chisq_bs_fft(vis_arr, A, bis, sigma):
    """Bispectrum chi-squared from fft"""

    im_info, sampler_info_list, gridder_info_list = A
    bisamples = obsh.sampler(vis_arr, sampler_info_list, sample_type="bs")

    return np.sum(np.abs(((bis - bisamples)/sigma))**2)/(2.*len(bis))


def chisqgrad_bs_fft(vis_arr, A, bis, sigma):
    """The gradient of the amplitude chi-squared
    """
    im_info, sampler_info_list, gridder_info_list = A

    v1 = obsh.sampler(vis_arr, [sampler_info_list[0]], sample_type="vis")
    v2 = obsh.sampler(vis_arr, [sampler_info_list[1]], sample_type="vis")
    v3 = obsh.sampler(vis_arr, [sampler_info_list[2]], sample_type="vis")
    bisamples = v1*v2*v3

    wdiff = -1.0/len(bis)*(bis - bisamples)/(sigma**2)

    pt1 = wdiff * (v2 * v3).conj() * sampler_info_list[0].pulsefac.conj()
    pt2 = wdiff * (v1 * v3).conj() * sampler_info_list[1].pulsefac.conj()
    pt3 = wdiff * (v1 * v2).conj() * sampler_info_list[2].pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff = obsh.gridder([pt1, pt2, pt3], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevant cells and flatten
    # TODO or is x<-->y??
    out = np.real(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                           im_info.padvaly1:-im_info.padvaly2].flatten())
    return out


def chisq_cphase_fft(vis_arr, A, clphase, sigma):
    """Closure Phases (normalized) chi-squared from fft
    """

    clphase = clphase * ehc.DEGREE
    sigma = sigma * ehc.DEGREE

    im_info, sampler_info_list, gridder_info_list = A
    clphase_samples = np.angle(obsh.sampler(vis_arr, sampler_info_list, sample_type="bs"))

    chisq = (2.0/len(clphase)) * np.sum((1.0 - np.cos(clphase-clphase_samples))/(sigma**2))
    return chisq


def chisqgrad_cphase_fft(vis_arr, A, clphase, sigma):
    """The gradient of the closure phase chi-squared from fft"""

    clphase = clphase * ehc.DEGREE
    sigma = sigma * ehc.DEGREE
    im_info, sampler_info_list, gridder_info_list = A

    # sample visibilities and closure phases
    v1 = obsh.sampler(vis_arr, [sampler_info_list[0]], sample_type="vis")
    v2 = obsh.sampler(vis_arr, [sampler_info_list[1]], sample_type="vis")
    v3 = obsh.sampler(vis_arr, [sampler_info_list[2]], sample_type="vis")
    clphase_samples = np.angle(v1*v2*v3)

    pref = (2.0/len(clphase)) * np.sin(clphase - clphase_samples)/(sigma**2)
    pt1 = pref/v1.conj() * sampler_info_list[0].pulsefac.conj()
    pt2 = pref/v2.conj() * sampler_info_list[1].pulsefac.conj()
    pt3 = pref/v3.conj() * sampler_info_list[2].pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff = obsh.gridder([pt1, pt2, pt3], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevant cells and flatten
    # TODO or is x<-->y??
    out = np.imag(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                           im_info.padvaly1:-im_info.padvaly2].flatten())

    return out


def chisq_cphase_diag_fft(imvec, A, clphase_diag, sigma):
    """Diagonalized closure phases (normalized) chi-squared from fft
    """

    clphase_diag = np.concatenate(clphase_diag) * ehc.DEGREE
    sigma = np.concatenate(sigma) * ehc.DEGREE

    A3 = A[0]
    tform_mats = A[1]

    im_info, sampler_info_list, gridder_info_list = A3
    vis_arr = obsh.fft_imvec(imvec, A3[0])
    clphase_samples = np.angle(obsh.sampler(vis_arr, sampler_info_list, sample_type="bs"))

    count = 0
    clphase_diag_samples = []
    for tform_mat in tform_mats:
        clphase_samples_here = clphase_samples[count:count+len(tform_mat)]
        clphase_diag_samples.append(np.dot(tform_mat, clphase_samples_here))
        count += len(tform_mat)

    clphase_diag_samples = np.concatenate(clphase_diag_samples)

    chisq = np.sum((1.0 - np.cos(clphase_diag-clphase_diag_samples))/(sigma**2))
    chisq *= (2.0/len(clphase_diag))
    return chisq


def chisqgrad_cphase_diag_fft(imvec, A, clphase_diag, sigma):
    """The gradient of the closure phase chi-squared from fft"""

    clphase_diag = np.concatenate(clphase_diag) * ehc.DEGREE
    sigma = np.concatenate(sigma) * ehc.DEGREE

    A3 = A[0]
    tform_mats = A[1]

    im_info, sampler_info_list, gridder_info_list = A3
    vis_arr = obsh.fft_imvec(imvec, A3[0])

    # sample visibilities and closure phases
    v1 = obsh.sampler(vis_arr, [sampler_info_list[0]], sample_type="vis")
    v2 = obsh.sampler(vis_arr, [sampler_info_list[1]], sample_type="vis")
    v3 = obsh.sampler(vis_arr, [sampler_info_list[2]], sample_type="vis")
    clphase_samples = np.angle(v1*v2*v3)

    # gradient vec stuff
    count = 0
    pref = np.zeros_like(clphase_samples)
    for tform_mat in tform_mats:

        clphase_diag_samples = np.dot(tform_mat, clphase_samples[count:count+len(tform_mat)])
        clphase_diag_measured = clphase_diag[count:count+len(tform_mat)]
        clphase_diag_sigma = sigma[count:count+len(tform_mat)]

        for j in range(len(clphase_diag_measured)):
            pref[count:count+len(tform_mat)] += 2.0 * tform_mat[j, :] * np.sin(
                clphase_diag_measured[j] - clphase_diag_samples[j])/(clphase_diag_sigma[j]**2)

        count += len(tform_mat)

    pt1 = pref/v1.conj() * sampler_info_list[0].pulsefac.conj()
    pt2 = pref/v2.conj() * sampler_info_list[1].pulsefac.conj()
    pt3 = pref/v3.conj() * sampler_info_list[2].pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff = obsh.gridder([pt1, pt2, pt3], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevant cells and flatten
    deriv = np.imag(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                             im_info.padvaly1:-im_info.padvaly2].flatten())
    deriv *= 1.0/np.float(len(clphase_diag))

    return deriv


def chisq_camp_fft(vis_arr, A, clamp, sigma):
    """Closure Amplitudes (normalized) chi-squared from fft
    """

    im_info, sampler_info_list, gridder_info_list = A
    clamp_samples = obsh.sampler(vis_arr, sampler_info_list, sample_type="camp")
    chisq = np.sum(np.abs((clamp - clamp_samples)/sigma)**2)/len(clamp)
    return chisq


def chisqgrad_camp_fft(vis_arr, A, clamp, sigma):
    """The gradient of the closure amplitude chi-squared from fft
    """

    im_info, sampler_info_list, gridder_info_list = A

    # sampled visibility and closure amplitudes
    v1 = obsh.sampler(vis_arr, [sampler_info_list[0]], sample_type="vis")
    v2 = obsh.sampler(vis_arr, [sampler_info_list[1]], sample_type="vis")
    v3 = obsh.sampler(vis_arr, [sampler_info_list[2]], sample_type="vis")
    v4 = obsh.sampler(vis_arr, [sampler_info_list[3]], sample_type="vis")
    clamp_samples = np.abs((v1 * v2)/(v3 * v4))

    # gradient components
    pp = (-2.0/len(clamp)) * ((clamp - clamp_samples) * clamp_samples)/(sigma**2)
    pt1 = pp/v1.conj() * sampler_info_list[0].pulsefac.conj()
    pt2 = pp/v2.conj() * sampler_info_list[1].pulsefac.conj()
    pt3 = -pp/v3.conj() * sampler_info_list[2].pulsefac.conj()
    pt4 = -pp/v4.conj() * sampler_info_list[3].pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff = obsh.gridder([pt1, pt2, pt3, pt4], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevant cells and flatten
    # TODO or is x<-->y??
    out = np.real(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                           im_info.padvaly1:-im_info.padvaly2].flatten())

    return out


def chisq_logcamp_fft(vis_arr, A, log_clamp, sigma):
    """Closure Amplitudes (normalized) chi-squared from fft
    """

    im_info, sampler_info_list, gridder_info_list = A
    log_clamp_samples = np.log(obsh.sampler(vis_arr, sampler_info_list, sample_type='camp'))

    chisq = np.sum(np.abs((log_clamp - log_clamp_samples)/sigma)**2) / (len(log_clamp))

    return chisq


def chisqgrad_logcamp_fft(vis_arr, A, log_clamp, sigma):
    """The gradient of the closure amplitude chi-squared from fft
    """

    im_info, sampler_info_list, gridder_info_list = A

    # sampled visibility and closure amplitudes
    v1 = obsh.sampler(vis_arr, [sampler_info_list[0]], sample_type="vis")
    v2 = obsh.sampler(vis_arr, [sampler_info_list[1]], sample_type="vis")
    v3 = obsh.sampler(vis_arr, [sampler_info_list[2]], sample_type="vis")
    v4 = obsh.sampler(vis_arr, [sampler_info_list[3]], sample_type="vis")

    log_clamp_samples = np.log(np.abs((v1 * v2)/(v3 * v4)))

    # gradient components
    pp = (-2.0/len(log_clamp)) * (log_clamp - log_clamp_samples) / (sigma**2)
    pt1 = pp / v1.conj() * sampler_info_list[0].pulsefac.conj()
    pt2 = pp / v2.conj() * sampler_info_list[1].pulsefac.conj()
    pt3 = -pp / v3.conj() * sampler_info_list[2].pulsefac.conj()
    pt4 = -pp / v4.conj() * sampler_info_list[3].pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff = obsh.gridder([pt1, pt2, pt3, pt4], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevant cells and flatten
    # TODO or is x<-->y??
    out = np.real(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                           im_info.padvaly1:-im_info.padvaly2].flatten())

    return out


def chisq_logcamp_diag_fft(imvec, A, log_clamp_diag, sigma):
    """Diagonalized log closure amplitudes (normalized) chi-squared from fft
    """

    log_clamp_diag = np.concatenate(log_clamp_diag)
    sigma = np.concatenate(sigma)

    A4 = A[0]
    tform_mats = A[1]

    im_info, sampler_info_list, gridder_info_list = A4
    vis_arr = obsh.fft_imvec(imvec, A4[0])
    log_clamp_samples = np.log(obsh.sampler(vis_arr, sampler_info_list, sample_type='camp'))

    count = 0
    log_clamp_diag_samples = []
    for tform_mat in tform_mats:
        log_clamp_samples_here = log_clamp_samples[count:count+len(tform_mat)]
        log_clamp_diag_samples.append(np.dot(tform_mat, log_clamp_samples_here))
        count += len(tform_mat)
    log_clamp_diag_samples = np.concatenate(log_clamp_diag_samples)

    chisq = np.sum(np.abs((log_clamp_diag - log_clamp_diag_samples)/sigma)**2)
    chisq /= (len(log_clamp_diag))
    return chisq


def chisqgrad_logcamp_diag_fft(imvec, A, log_clamp_diag, sigma):
    """The gradient of the diagonalized log closure amplitude chi-squared from fft
    """

    log_clamp_diag = np.concatenate(log_clamp_diag)
    sigma = np.concatenate(sigma)

    A4 = A[0]
    tform_mats = A[1]

    im_info, sampler_info_list, gridder_info_list = A4
    vis_arr = obsh.fft_imvec(imvec, A4[0])

    # sampled visibility and closure amplitudes
    v1 = obsh.sampler(vis_arr, [sampler_info_list[0]], sample_type="vis")
    v2 = obsh.sampler(vis_arr, [sampler_info_list[1]], sample_type="vis")
    v3 = obsh.sampler(vis_arr, [sampler_info_list[2]], sample_type="vis")
    v4 = obsh.sampler(vis_arr, [sampler_info_list[3]], sample_type="vis")
    log_clamp_samples = np.log(np.abs((v1 * v2)/(v3 * v4)))

    # gradient vec stuff
    count = 0
    pref = np.zeros_like(log_clamp_samples)
    for tform_mat in tform_mats:

        log_clamp_diag_samples = np.dot(tform_mat, log_clamp_samples[count:count+len(tform_mat)])
        log_clamp_diag_measured = log_clamp_diag[count:count+len(tform_mat)]
        log_clamp_diag_sigma = sigma[count:count+len(tform_mat)]

        for j in range(len(log_clamp_diag_measured)):
            pref[count:count+len(tform_mat)] += -2.0 * tform_mat[j, :] * \
                (log_clamp_diag_measured[j] - log_clamp_diag_samples[j]) / \
                (log_clamp_diag_sigma[j]**2)

        count += len(tform_mat)

    pt1 = pref / v1.conj() * sampler_info_list[0].pulsefac.conj()
    pt2 = pref / v2.conj() * sampler_info_list[1].pulsefac.conj()
    pt3 = -pref / v3.conj() * sampler_info_list[2].pulsefac.conj()
    pt4 = -pref / v4.conj() * sampler_info_list[3].pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff = obsh.gridder([pt1, pt2, pt3, pt4], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevant cells and flatten
    deriv = np.real(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                             im_info.padvaly1:-im_info.padvaly2].flatten())
    deriv *= 1.0/np.float(len(log_clamp_diag))

    return deriv


def chisq_logamp_fft(vis_arr, A, amp, sigma):
    """Visibility amplitude chi-squared from fft
    """

    # to lowest order the variance on the logarithm of a quantity x is
    # sigma^2_log(x) = sigma^2/x^2
    logsigma = sigma / amp

    im_info, sampler_info_list, gridder_info_list = A
    amp_samples = np.abs(obsh.sampler(vis_arr, sampler_info_list, sample_type="vis"))
    chisq = np.sum(np.abs((np.log(amp_samples)-np.log(amp))/logsigma)**2)/(len(amp))
    return chisq


def chisqgrad_logamp_fft(vis_arr, A, amp, sigma):
    """The gradient of the amplitude chi-kernesquared
    """

    im_info, sampler_info_list, gridder_info_list = A

    # samples
    samples = obsh.sampler(vis_arr, sampler_info_list, sample_type="vis")
    amp_samples = np.abs(samples)

    # gradient FT
    logsigma = sigma / amp
    pulsefac = sampler_info_list[0].pulsefac
    wdiff_vec = (-2.0/len(amp)*((np.log(amp) - np.log(amp_samples))) /
                 (logsigma**2) / samples.conj()) * pulsefac.conj()

    # Setup and perform the inverse FFT
    wdiff_arr = obsh.gridder([wdiff_vec], gridder_info_list)
    grad_arr = np.fft.ifftshift(np.fft.ifft2(np.fft.fftshift(wdiff_arr)))
    grad_arr = grad_arr * (im_info.npad * im_info.npad)

    # extract relevent cells and flatten
    # TODO or is x<-->y??
    out = np.real(grad_arr[im_info.padvalx1:-im_info.padvalx2,
                           im_info.padvaly1:-im_info.padvaly2].flatten())

    return out

##################################################################################################
# NFFT Chi-squared and Gradient Functions
##################################################################################################


def chisq_vis_nfft(imvec, A, vis, sigma):
    """Visibility chi-squared from nfft
    """

    # get nfft object
    nfft_info = A[0]
    plan = nfft_info.plan
    pulsefac = nfft_info.pulsefac

    # compute uniform --> nonuniform transform
    plan.f_hat = imvec.copy().reshape((nfft_info.ydim, nfft_info.xdim)).T
    plan.trafo()
    samples = plan.f.copy()*pulsefac

    # compute chi^2
    chisq = np.sum(np.abs((samples-vis)/sigma)**2)/(2*len(vis))

    return chisq


def chisqgrad_vis_nfft(imvec, A, vis, sigma):
    """The gradient of the visibility chi-squared from nfft
    """

    # get nfft object
    nfft_info = A[0]
    plan = nfft_info.plan
    pulsefac = nfft_info.pulsefac

    # compute uniform --> nonuniform transform
    plan.f_hat = imvec.copy().reshape((nfft_info.ydim, nfft_info.xdim)).T
    plan.trafo()
    samples = plan.f.copy()*pulsefac

    # gradient vec for adjoint FT
    wdiff_vec = (-1.0/len(vis)*(vis - samples)/(sigma**2)) * pulsefac.conj()
    plan.f = wdiff_vec
    plan.adjoint()
    grad = np.real((plan.f_hat.copy().T).reshape(nfft_info.xdim*nfft_info.ydim))

    return grad


def chisq_amp_nfft(imvec, A, amp, sigma):
    """Visibility amplitude chi-squared from nfft
    """
    # get nfft object
    nfft_info = A[0]
    plan = nfft_info.plan
    pulsefac = nfft_info.pulsefac

    # compute uniform --> nonuniform transform
    plan.f_hat = imvec.copy().reshape((nfft_info.ydim, nfft_info.xdim)).T
    plan.trafo()
    samples = plan.f.copy()*pulsefac

    # compute chi^2
    amp_samples = np.abs(samples)
    chisq = np.sum(np.abs((amp_samples-amp)/sigma)**2)/(len(amp))

    return chisq


def chisqgrad_amp_nfft(imvec, A, amp, sigma):
    """The gradient of the amplitude chi-squared from nfft
    """

    # get nfft object
    nfft_info = A[0]
    plan = nfft_info.plan
    pulsefac = nfft_info.pulsefac

    # compute uniform --> nonuniform transform
    plan.f_hat = imvec.copy().reshape((nfft_info.ydim, nfft_info.xdim)).T
    plan.trafo()
    samples = plan.f.copy()*pulsefac
    amp_samples = np.abs(samples)

    # gradient vec for adjoint FT
    wdiff_vec = (-2.0/len(amp)*((amp - amp_samples) * samples) /
                 (sigma**2) / amp_samples) * pulsefac.conj()
    plan.f = wdiff_vec
    plan.adjoint()
    out = np.real((plan.f_hat.copy().T).reshape(nfft_info.xdim*nfft_info.ydim))

    return out


def chisq_bs_nfft(imvec, A, bis, sigma):
    """Bispectrum chi-squared from fft"""

    # get nfft objects
    nfft_info1 = A[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    samples1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    samples2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    samples3 = plan3.f.copy()*pulsefac3

    # compute chi^2
    bisamples = samples1*samples2*samples3
    chisq = np.sum(np.abs(((bis - bisamples)/sigma))**2)/(2.*len(bis))
    return chisq


def chisqgrad_bs_nfft(imvec, A, bis, sigma):
    """The gradient of the amplitude chi-squared from the nfft
    """
    # get nfft objects
    nfft_info1 = A[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    v1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    v2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    v3 = plan3.f.copy()*pulsefac3

    # gradient vec for adjoint FT
    bisamples = v1*v2*v3
    wdiff = -1.0/len(bis)*(bis - bisamples)/(sigma**2)
    pt1 = wdiff * (v2 * v3).conj() * pulsefac1.conj()
    pt2 = wdiff * (v1 * v3).conj() * pulsefac2.conj()
    pt3 = wdiff * (v1 * v2).conj() * pulsefac3.conj()

    # Setup and perform the inverse FFT
    plan1.f = pt1
    plan1.adjoint()
    out1 = np.real((plan1.f_hat.copy().T).reshape(nfft_info1.xdim*nfft_info1.ydim))

    plan2.f = pt2
    plan2.adjoint()
    out2 = np.real((plan2.f_hat.copy().T).reshape(nfft_info2.xdim*nfft_info2.ydim))

    plan3.f = pt3
    plan3.adjoint()
    out3 = np.real((plan3.f_hat.copy().T).reshape(nfft_info3.xdim*nfft_info3.ydim))

    out = out1 + out2 + out3
    return out


def chisq_cphase_nfft(imvec, A, clphase, sigma):
    """Closure Phases (normalized) chi-squared from nfft
    """

    clphase = clphase * ehc.DEGREE
    sigma = sigma * ehc.DEGREE

    # get nfft objects
    nfft_info1 = A[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    samples1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    samples2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    samples3 = plan3.f.copy()*pulsefac3

    # compute chi^2
    clphase_samples = np.angle(samples1*samples2*samples3)
    chisq = (2.0/len(clphase)) * np.sum((1.0 - np.cos(clphase-clphase_samples))/(sigma**2))

    return chisq


def chisqgrad_cphase_nfft(imvec, A, clphase, sigma):
    """The gradient of the closure phase chi-squared from nfft"""

    clphase = clphase * ehc.DEGREE
    sigma = sigma * ehc.DEGREE

    # get nfft objects
    nfft_info1 = A[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    v1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    v2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    v3 = plan3.f.copy()*pulsefac3

    # gradient vec for adjoint FT
    clphase_samples = np.angle(v1*v2*v3)
    pref = (2.0/len(clphase)) * np.sin(clphase - clphase_samples)/(sigma**2)
    pt1 = pref/v1.conj() * pulsefac1.conj()
    pt2 = pref/v2.conj() * pulsefac2.conj()
    pt3 = pref/v3.conj() * pulsefac3.conj()

    # Setup and perform the inverse FFT
    plan1.f = pt1
    plan1.adjoint()
    out1 = np.imag((plan1.f_hat.copy().T).reshape(nfft_info1.xdim*nfft_info1.ydim))

    plan2.f = pt2
    plan2.adjoint()
    out2 = np.imag((plan2.f_hat.copy().T).reshape(nfft_info2.xdim*nfft_info2.ydim))

    plan3.f = pt3
    plan3.adjoint()
    out3 = np.imag((plan3.f_hat.copy().T).reshape(nfft_info3.xdim*nfft_info3.ydim))

    out = out1 + out2 + out3
    return out


def chisq_cphase_diag_nfft(imvec, A, clphase_diag, sigma):
    """Diagonalized closure phases (normalized) chi-squared from nfft
    """

    clphase_diag = np.concatenate(clphase_diag) * ehc.DEGREE
    sigma = np.concatenate(sigma) * ehc.DEGREE

    A3 = A[0]
    tform_mats = A[1]

    # get nfft objects
    nfft_info1 = A3[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A3[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A3[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    samples1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    samples2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    samples3 = plan3.f.copy()*pulsefac3

    clphase_samples = np.angle(samples1*samples2*samples3)

    count = 0
    clphase_diag_samples = []
    for tform_mat in tform_mats:
        clphase_samples_here = clphase_samples[count:count+len(tform_mat)]
        clphase_diag_samples.append(np.dot(tform_mat, clphase_samples_here))
        count += len(tform_mat)

    clphase_diag_samples = np.concatenate(clphase_diag_samples)

    # compute chi^2
    chisq = (2.0/len(clphase_diag)) * \
        np.sum((1.0 - np.cos(clphase_diag-clphase_diag_samples))/(sigma**2))

    return chisq


def chisqgrad_cphase_diag_nfft(imvec, A, clphase_diag, sigma):
    """The gradient of the diagonalized closure phase chi-squared from nfft"""

    clphase_diag = np.concatenate(clphase_diag) * ehc.DEGREE
    sigma = np.concatenate(sigma) * ehc.DEGREE

    A3 = A[0]
    tform_mats = A[1]

    # get nfft objects
    nfft_info1 = A3[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A3[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A3[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    v1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    v2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    v3 = plan3.f.copy()*pulsefac3

    clphase_samples = np.angle(v1*v2*v3)

    # gradient vec for adjoint FT
    count = 0
    pref = np.zeros_like(clphase_samples)
    for tform_mat in tform_mats:

        clphase_diag_samples = np.dot(tform_mat, clphase_samples[count:count+len(tform_mat)])
        clphase_diag_measured = clphase_diag[count:count+len(tform_mat)]
        clphase_diag_sigma = sigma[count:count+len(tform_mat)]

        for j in range(len(clphase_diag_measured)):
            pref[count:count+len(tform_mat)] += 2.0 * tform_mat[j, :] * np.sin(
                clphase_diag_measured[j] - clphase_diag_samples[j])/(clphase_diag_sigma[j]**2)

        count += len(tform_mat)

    pt1 = pref/v1.conj() * pulsefac1.conj()
    pt2 = pref/v2.conj() * pulsefac2.conj()
    pt3 = pref/v3.conj() * pulsefac3.conj()

    # Setup and perform the inverse FFT
    plan1.f = pt1
    plan1.adjoint()
    out1 = np.imag((plan1.f_hat.copy().T).reshape(nfft_info1.xdim*nfft_info1.ydim))

    plan2.f = pt2
    plan2.adjoint()
    out2 = np.imag((plan2.f_hat.copy().T).reshape(nfft_info2.xdim*nfft_info2.ydim))

    plan3.f = pt3
    plan3.adjoint()
    out3 = np.imag((plan3.f_hat.copy().T).reshape(nfft_info3.xdim*nfft_info3.ydim))

    deriv = out1 + out2 + out3
    deriv *= 1.0/np.float(len(clphase_diag))

    return deriv


def chisq_camp_nfft(imvec, A, clamp, sigma):
    """Closure Amplitudes (normalized) chi-squared from fft
    """

    # get nfft objects
    nfft_info1 = A[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    nfft_info4 = A[3]
    plan4 = nfft_info4.plan
    pulsefac4 = nfft_info4.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    samples1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    samples2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    samples3 = plan3.f.copy()*pulsefac3

    plan4.f_hat = imvec.copy().reshape((nfft_info4.ydim, nfft_info4.xdim)).T
    plan4.trafo()
    samples4 = plan4.f.copy()*pulsefac4

    # compute chi^2
    clamp_samples = np.abs((samples1*samples2)/(samples3*samples4))
    chisq = np.sum(np.abs((clamp - clamp_samples)/sigma)**2)/len(clamp)
    return chisq


def chisqgrad_camp_nfft(imvec, A, clamp, sigma):
    """The gradient of the closure amplitude chi-squared from fft
    """

    # get nfft objects
    nfft_info1 = A[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    nfft_info4 = A[3]
    plan4 = nfft_info4.plan
    pulsefac4 = nfft_info4.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    v1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    v2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    v3 = plan3.f.copy()*pulsefac3

    plan4.f_hat = imvec.copy().reshape((nfft_info4.ydim, nfft_info4.xdim)).T
    plan4.trafo()
    v4 = plan4.f.copy()*pulsefac4

    # gradient vec for adjoint FT
    clamp_samples = np.abs((v1 * v2)/(v3 * v4))

    pp = (-2.0/len(clamp)) * ((clamp - clamp_samples) * clamp_samples)/(sigma**2)
    pt1 = pp/v1.conj() * pulsefac1.conj()
    pt2 = pp/v2.conj() * pulsefac2.conj()
    pt3 = -pp/v3.conj() * pulsefac3.conj()
    pt4 = -pp/v4.conj() * pulsefac4.conj()

    # Setup and perform the inverse FFT
    plan1.f = pt1
    plan1.adjoint()
    out1 = np.real((plan1.f_hat.copy().T).reshape(nfft_info1.xdim*nfft_info1.ydim))

    plan2.f = pt2
    plan2.adjoint()
    out2 = np.real((plan2.f_hat.copy().T).reshape(nfft_info2.xdim*nfft_info2.ydim))

    plan3.f = pt3
    plan3.adjoint()
    out3 = np.real((plan3.f_hat.copy().T).reshape(nfft_info3.xdim*nfft_info3.ydim))

    plan4.f = pt4
    plan4.adjoint()
    out4 = np.real((plan4.f_hat.copy().T).reshape(nfft_info4.xdim*nfft_info4.ydim))

    out = out1 + out2 + out3 + out4
    return out


def chisq_logcamp_nfft(imvec, A, log_clamp, sigma):
    """Log Closure Amplitudes (normalized) chi-squared from fft
    """

    # get nfft objects
    nfft_info1 = A[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    nfft_info4 = A[3]
    plan4 = nfft_info4.plan
    pulsefac4 = nfft_info4.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    samples1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    samples2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    samples3 = plan3.f.copy()*pulsefac3

    plan4.f_hat = imvec.copy().reshape((nfft_info4.ydim, nfft_info4.xdim)).T
    plan4.trafo()
    samples4 = plan4.f.copy()*pulsefac4

    # compute chi^2
    log_clamp_samples = (np.log(np.abs(samples1)) + np.log(np.abs(samples2)) -
                         np.log(np.abs(samples3)) - np.log(np.abs(samples4)))
    chisq = np.sum(np.abs((log_clamp - log_clamp_samples)/sigma)**2) / (len(log_clamp))
    return chisq


def chisqgrad_logcamp_nfft(imvec, A, log_clamp, sigma):
    """The gradient of the log closure amplitude chi-squared from fft
    """

    # get nfft objects
    nfft_info1 = A[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    nfft_info4 = A[3]
    plan4 = nfft_info4.plan
    pulsefac4 = nfft_info4.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    v1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    v2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    v3 = plan3.f.copy()*pulsefac3

    plan4.f_hat = imvec.copy().reshape((nfft_info4.ydim, nfft_info4.xdim)).T
    plan4.trafo()
    v4 = plan4.f.copy()*pulsefac4

    # gradient vec for adjoint FT
    log_clamp_samples = np.log(np.abs((v1 * v2)/(v3 * v4)))

    pp = (-2.0/len(log_clamp)) * (log_clamp - log_clamp_samples) / (sigma**2)
    pt1 = pp / v1.conj() * pulsefac1.conj()
    pt2 = pp / v2.conj() * pulsefac2.conj()
    pt3 = -pp / v3.conj() * pulsefac3.conj()
    pt4 = -pp / v4.conj() * pulsefac4.conj()

    # Setup and perform the inverse FFT
    plan1.f = pt1
    plan1.adjoint()
    out1 = np.real((plan1.f_hat.copy().T).reshape(nfft_info1.xdim*nfft_info1.ydim))

    plan2.f = pt2
    plan2.adjoint()
    out2 = np.real((plan2.f_hat.copy().T).reshape(nfft_info2.xdim*nfft_info2.ydim))

    plan3.f = pt3
    plan3.adjoint()
    out3 = np.real((plan3.f_hat.copy().T).reshape(nfft_info3.xdim*nfft_info3.ydim))

    plan4.f = pt4
    plan4.adjoint()
    out4 = np.real((plan4.f_hat.copy().T).reshape(nfft_info4.xdim*nfft_info4.ydim))

    out = out1 + out2 + out3 + out4
    return out


def chisq_logcamp_diag_nfft(imvec, A, log_clamp_diag, sigma):
    """Diagonalized log closure amplitudes (normalized) chi-squared from nfft
    """

    log_clamp_diag = np.concatenate(log_clamp_diag)
    sigma = np.concatenate(sigma)

    A4 = A[0]
    tform_mats = A[1]

    # get nfft objects
    nfft_info1 = A4[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A4[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A4[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    nfft_info4 = A4[3]
    plan4 = nfft_info4.plan
    pulsefac4 = nfft_info4.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    samples1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    samples2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    samples3 = plan3.f.copy()*pulsefac3

    plan4.f_hat = imvec.copy().reshape((nfft_info4.ydim, nfft_info4.xdim)).T
    plan4.trafo()
    samples4 = plan4.f.copy()*pulsefac4

    log_clamp_samples = (np.log(np.abs(samples1)) + np.log(np.abs(samples2)) -
                         np.log(np.abs(samples3)) - np.log(np.abs(samples4)))

    count = 0
    log_clamp_diag_samples = []
    for tform_mat in tform_mats:
        log_clamp_samples_here = log_clamp_samples[count:count+len(tform_mat)]
        log_clamp_diag_samples.append(np.dot(tform_mat, log_clamp_samples_here))
        count += len(tform_mat)

    log_clamp_diag_samples = np.concatenate(log_clamp_diag_samples)

    # compute chi^2
    chisq = np.sum(np.abs((log_clamp_diag - log_clamp_diag_samples)/sigma)**2) / \
        (len(log_clamp_diag))

    return chisq


def chisqgrad_logcamp_diag_nfft(imvec, A, log_clamp_diag, sigma):
    """The gradient of the diagonalized log closure amplitude chi-squared from fft
    """

    log_clamp_diag = np.concatenate(log_clamp_diag)
    sigma = np.concatenate(sigma)

    A4 = A[0]
    tform_mats = A[1]

    # get nfft objects
    nfft_info1 = A4[0]
    plan1 = nfft_info1.plan
    pulsefac1 = nfft_info1.pulsefac

    nfft_info2 = A4[1]
    plan2 = nfft_info2.plan
    pulsefac2 = nfft_info2.pulsefac

    nfft_info3 = A4[2]
    plan3 = nfft_info3.plan
    pulsefac3 = nfft_info3.pulsefac

    nfft_info4 = A4[3]
    plan4 = nfft_info4.plan
    pulsefac4 = nfft_info4.pulsefac

    # compute uniform --> nonuniform transforms
    plan1.f_hat = imvec.copy().reshape((nfft_info1.ydim, nfft_info1.xdim)).T
    plan1.trafo()
    v1 = plan1.f.copy()*pulsefac1

    plan2.f_hat = imvec.copy().reshape((nfft_info2.ydim, nfft_info2.xdim)).T
    plan2.trafo()
    v2 = plan2.f.copy()*pulsefac2

    plan3.f_hat = imvec.copy().reshape((nfft_info3.ydim, nfft_info3.xdim)).T
    plan3.trafo()
    v3 = plan3.f.copy()*pulsefac3

    plan4.f_hat = imvec.copy().reshape((nfft_info4.ydim, nfft_info4.xdim)).T
    plan4.trafo()
    v4 = plan4.f.copy()*pulsefac4

    log_clamp_samples = np.log(np.abs((v1 * v2)/(v3 * v4)))

    # gradient vec for adjoint FT
    count = 0
    pp = np.zeros_like(log_clamp_samples)
    for tform_mat in tform_mats:

        log_clamp_diag_samples = np.dot(tform_mat, log_clamp_samples[count:count+len(tform_mat)])
        log_clamp_diag_measured = log_clamp_diag[count:count+len(tform_mat)]
        log_clamp_diag_sigma = sigma[count:count+len(tform_mat)]

        for j in range(len(log_clamp_diag_measured)):
            pp[count:count+len(tform_mat)] += -2.0 * tform_mat[j, :] * \
                (log_clamp_diag_measured[j] - log_clamp_diag_samples[j]) / \
                (log_clamp_diag_sigma[j]**2)

        count += len(tform_mat)

    pt1 = pp / v1.conj() * pulsefac1.conj()
    pt2 = pp / v2.conj() * pulsefac2.conj()
    pt3 = -pp / v3.conj() * pulsefac3.conj()
    pt4 = -pp / v4.conj() * pulsefac4.conj()

    # Setup and perform the inverse FFT
    plan1.f = pt1
    plan1.adjoint()
    out1 = np.real((plan1.f_hat.copy().T).reshape(nfft_info1.xdim*nfft_info1.ydim))

    plan2.f = pt2
    plan2.adjoint()
    out2 = np.real((plan2.f_hat.copy().T).reshape(nfft_info2.xdim*nfft_info2.ydim))

    plan3.f = pt3
    plan3.adjoint()
    out3 = np.real((plan3.f_hat.copy().T).reshape(nfft_info3.xdim*nfft_info3.ydim))

    plan4.f = pt4
    plan4.adjoint()
    out4 = np.real((plan4.f_hat.copy().T).reshape(nfft_info4.xdim*nfft_info4.ydim))

    deriv = out1 + out2 + out3 + out4
    deriv *= 1.0/np.float(len(log_clamp_diag))

    return deriv


def chisq_logamp_nfft(imvec, A, amp, sigma):
    """Visibility log amplitude chi-squared from nfft
    """

    # get nfft object
    nfft_info = A[0]
    plan = nfft_info.plan
    pulsefac = nfft_info.pulsefac

    # compute uniform --> nonuniform transform
    plan.f_hat = imvec.copy().reshape((nfft_info.ydim, nfft_info.xdim)).T
    plan.trafo()
    samples = plan.f.copy()*pulsefac

    # to lowest order the variance on the logarithm of a quantity x is
    # sigma^2_log(x) = sigma^2/x^2
    logsigma = sigma / amp

    # compute chi^2
    amp_samples = np.abs(samples)
    chisq = np.sum(np.abs((np.log(amp_samples)-np.log(amp))/logsigma)**2)/(len(amp))

    return chisq


def chisqgrad_logamp_nfft(imvec, A, amp, sigma):
    """The gradient of the log amplitude chi-squared from nfft
    """

    # get nfft object
    nfft_info = A[0]
    plan = nfft_info.plan
    pulsefac = nfft_info.pulsefac

    # compute uniform --> nonuniform transform
    plan.f_hat = imvec.copy().reshape((nfft_info.ydim, nfft_info.xdim)).T
    plan.trafo()
    samples = plan.f.copy()*pulsefac
    amp_samples = np.abs(samples)

    # gradient vec for adjoint FT
    logsigma = sigma / amp
    wdiff_vec = (-2.0/len(amp)*((np.log(amp) - np.log(amp_samples))) /
                 (logsigma**2) / samples.conj()) * pulsefac.conj()
    plan.f = wdiff_vec
    plan.adjoint()
    out = np.real((plan.f_hat.copy().T).reshape(nfft_info.xdim*nfft_info.ydim))

    return out


##################################################################################################
# Regularizer and Gradient Functions
##################################################################################################

def sflux(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """Total flux constraint
    """
    if norm_reg:
        norm = flux**2
    else:
        norm = 1

    out = -(np.sum(imvec) - flux)**2
    return out/norm


def sfluxgrad(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """Total flux constraint gradient
    """
    if norm_reg:
        norm = flux**2
    else:
        norm = 1

    out = -2*(np.sum(imvec) - flux)*np.ones(len(imvec))
    return out / norm


def scm(imvec, nx, ny, psize, flux, embed_mask, norm_reg=NORM_REGULARIZER, beam_size=None):
    """Center-of-mass constraint
    """
    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = beam_size**2 * flux**2
    else:
        norm = 1

    xx, yy = np.meshgrid(range(nx//2, -nx//2, -1), range(ny//2, -ny//2, -1))
    xx = psize*xx.flatten()[embed_mask]
    yy = psize*yy.flatten()[embed_mask]

    out = -(np.sum(imvec*xx)**2 + np.sum(imvec*yy)**2)
    return out/norm


def scmgrad(imvec, nx, ny, psize, flux, embed_mask, norm_reg=NORM_REGULARIZER, beam_size=None):
    """Center-of-mass constraint gradient
    """
    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = beam_size**2 * flux**2
    else:
        norm = 1

    xx, yy = np.meshgrid(range(nx//2, -nx//2, -1), range(ny//2, -ny//2, -1))
    xx = psize*xx.flatten()[embed_mask]
    yy = psize*yy.flatten()[embed_mask]

    out = -2*(np.sum(imvec*xx)*xx + np.sum(imvec*yy)*yy)
    return out/norm


def ssimple(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """Simple entropy
    """
    if norm_reg:
        norm = flux
    else:
        norm = 1

    entropy = -np.sum(imvec*np.log(imvec/priorvec))
    return entropy/norm


def ssimplegrad(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """Simple entropy gradient
    """
    if norm_reg:
        norm = flux
    else:
        norm = 1

    entropygrad = -np.log(imvec/priorvec) - 1
    return entropygrad/norm


def sl1(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """L1 norm regularizer
    """
    if norm_reg:
        norm = flux
    else:
        norm = 1

    # l1 = -np.sum(np.abs(imvec - priorvec))
    l1 = -np.sum(np.abs(imvec))
    return l1/norm


def sl1grad(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """L1 norm gradient
    """
    if norm_reg:
        norm = flux
    else:
        norm = 1

    # l1grad = -np.sign(imvec - priorvec)
    l1grad = -np.sign(imvec)
    return l1grad/norm


def sl1w(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER, epsilon=ehc.EP):
    """Weighted L1 norm regularizer a la SMILI
    """

    if norm_reg:
        norm = 1  # should be ok?
        # This is SMILI normalization
        # norm = np.sum((np.sqrt(priorvec**2 + epsilon) + epsilon)/np.sqrt(priorvec**2 + epsilon))
    else:
        norm = 1

    num = np.sqrt(imvec**2 + epsilon)
    denom = np.sqrt(priorvec**2 + epsilon) + epsilon

    l1w = -np.sum(num/denom)
    return l1w/norm


def sl1wgrad(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER, epsilon=ehc.EP):
    """Weighted L1 norm gradient
    """
    if norm_reg:
        norm = 1  # should be ok?
        # This is SMILI normalization
        # norm = np.sum((np.sqrt(priorvec**2 + epsilon) + epsilon)/np.sqrt(priorvec**2 + epsilon))
    else:
        norm = 1

    num = imvec / np.sqrt(imvec**2 + epsilon)
    denom = np.sqrt(priorvec**2 + epsilon) + epsilon

    l1wgrad = - num / denom
    return l1wgrad/norm


def fA(imvec, I_ref=1.0, alpha_A=1.0):
    """Function to take imvec to itself in the limit alpha_A -> 0
       and to a binary representation in the limit alpha_A -> infinity
    """
    return 2.0/np.pi * (1.0 + alpha_A)/alpha_A * np.arctan(np.pi*alpha_A/2.0*np.abs(imvec)/I_ref)


def fAgrad(imvec, I_ref=1.0, alpha_A=1.0):
    """Function to take imvec to itself in the limit alpha_A -> 0
       and to a binary representation in the limit alpha_A -> infinity
    """
    return (1.0 + alpha_A) / (I_ref * (1.0 + (np.pi*alpha_A/2.0*imvec/I_ref)**2))


def slA(imvec, priorvec, psize, flux, beam_size=None, alpha_A=1.0, norm_reg=NORM_REGULARIZER):
    """l_A regularizer
    """

    # The appropriate I_ref is something like the total flux divided by the # of pixels per beam
    if beam_size is None:
        beam_size = psize
    I_ref = flux

    if norm_reg:
        norm_l1 = 1.0                  # as alpha_A ->0
        norm_l0 = (beam_size/psize)**2  # as alpha_A ->\infty
        weight_l1 = 1.0/(1.0 + alpha_A)
        weight_l0 = alpha_A
        norm = (norm_l1 * weight_l1 + norm_l0 * weight_l0)/(weight_l0 + weight_l1)
    else:
        norm = 1

    return -np.sum(fA(imvec, I_ref, alpha_A))/norm


def slAgrad(imvec, priorvec, psize, flux, beam_size=None, alpha_A=1.0, norm_reg=NORM_REGULARIZER):
    """l_A gradient
    """

    # The appropriate I_ref is something like the total flux divided by the # of pixels per beam
    if beam_size is None:
        beam_size = psize
    I_ref = flux

    if norm_reg:
        norm_l1 = 1.0                  # as alpha_A ->0
        norm_l0 = (beam_size/psize)**2  # as alpha_A ->\infty
        weight_l1 = 1.0/(1.0 + alpha_A)
        weight_l0 = alpha_A
        norm = (norm_l1 * weight_l1 + norm_l0 * weight_l0)/(weight_l0 + weight_l1)
    else:
        norm = 1

    return -fAgrad(imvec, I_ref, alpha_A)/norm


def sgs(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """Gull-skilling entropy
    """
    if norm_reg:
        norm = flux
    else:
        norm = 1

    entropy = np.sum(imvec - priorvec - imvec*np.log(imvec/priorvec))
    return entropy/norm


def sgsgrad(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """Gull-Skilling gradient
    """
    if norm_reg:
        norm = flux
    else:
        norm = 1

    entropygrad = -np.log(imvec/priorvec)
    return entropygrad/norm

# TODO: epsilon is 0 by default for backwards compatibilitys


def stv(imvec, nx, ny, psize, flux, norm_reg=NORM_REGULARIZER, beam_size=None, epsilon=0.):
    """Total variation regularizer
    """
    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = flux*psize / beam_size
    else:
        norm = 1

    im = imvec.reshape(ny, nx)
    impad = np.pad(im, 1, mode='constant', constant_values=0)
    im_l1 = np.roll(impad, -1, axis=0)[1:ny+1, 1:nx+1]
    im_l2 = np.roll(impad, -1, axis=1)[1:ny+1, 1:nx+1]
    out = -np.sum(np.sqrt(np.abs(im_l1 - im)**2 + np.abs(im_l2 - im)**2) + epsilon)
    return out/norm

# TODO: epsilon is 0 by default for backwards compatibility


def stvgrad(imvec, nx, ny, psize, flux, norm_reg=NORM_REGULARIZER, beam_size=None, epsilon=0.):
    """Total variation gradient
    """
    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = flux*psize / beam_size
    else:
        norm = 1

    im = imvec.reshape(ny, nx)
    impad = np.pad(im, 1, mode='constant', constant_values=0)
    im_l1 = np.roll(impad, -1, axis=0)[1:ny+1, 1:nx+1]
    im_l2 = np.roll(impad, -1, axis=1)[1:ny+1, 1:nx+1]
    im_r1 = np.roll(impad, 1, axis=0)[1:ny+1, 1:nx+1]
    im_r2 = np.roll(impad, 1, axis=1)[1:ny+1, 1:nx+1]

    # rotate images
    im_r1l2 = np.roll(np.roll(impad,  1, axis=0), -1, axis=1)[1:ny+1, 1:nx+1]
    im_l1r2 = np.roll(np.roll(impad, -1, axis=0), 1, axis=1)[1:ny+1, 1:nx+1]

    # add together terms and return
    g1 = (2*im - im_l1 - im_l2) / np.sqrt((im - im_l1)**2 + (im - im_l2)**2 + epsilon)
    g2 = (im - im_r1) / np.sqrt((im - im_r1)**2 + (im_r1l2 - im_r1)**2 + epsilon)
    g3 = (im - im_r2) / np.sqrt((im - im_r2)**2 + (im_l1r2 - im_r2)**2 + epsilon)

    # mask the first row column gradient terms that don't exist
    mask1 = np.zeros(im.shape)
    mask2 = np.zeros(im.shape)
    mask1[0, :] = 1
    mask2[:, 0] = 1
    g2[mask1.astype(bool)] = 0
    g3[mask2.astype(bool)] = 0

    # add terms together and return
    out = -(g1 + g2 + g3).flatten()
    return out/norm


def stv2(imvec, nx, ny, psize, flux, norm_reg=NORM_REGULARIZER, beam_size=None):
    """Squared Total variation regularizer
    """
    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = psize**4 * flux**2 / beam_size**4
    else:
        norm = 1

    im = imvec.reshape(ny, nx)
    impad = np.pad(im, 1, mode='constant', constant_values=0)
    im_l1 = np.roll(impad, -1, axis=0)[1:ny+1, 1:nx+1]
    im_l2 = np.roll(impad, -1, axis=1)[1:ny+1, 1:nx+1]
    out = -np.sum((im_l1 - im)**2 + (im_l2 - im)**2)
    return out/norm


def stv2grad(imvec, nx, ny, psize, flux, norm_reg=NORM_REGULARIZER, beam_size=None):
    """Squared Total variation gradient
    """
    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = psize**4 * flux**2 / beam_size**4
    else:
        norm = 1

    im = imvec.reshape(ny, nx)
    impad = np.pad(im, 1, mode='constant', constant_values=0)
    im_l1 = np.roll(impad, -1, axis=0)[1:ny+1, 1:nx+1]
    im_l2 = np.roll(impad, -1, axis=1)[1:ny+1, 1:nx+1]
    im_r1 = np.roll(impad, 1, axis=0)[1:ny+1, 1:nx+1]
    im_r2 = np.roll(impad, 1, axis=1)[1:ny+1, 1:nx+1]

    g1 = (2*im - im_l1 - im_l2)
    g2 = (im - im_r1)
    g3 = (im - im_r2)

    # mask the first row column gradient terms that don't exist
    mask1 = np.zeros(im.shape)
    mask2 = np.zeros(im.shape)
    mask1[0, :] = 1
    mask2[:, 0] = 1
    g2[mask1.astype(bool)] = 0
    g3[mask2.astype(bool)] = 0

    # add together terms and return
    out = -2*(g1 + g2 + g3).flatten()
    return out/norm


def spatch(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """Patch prior regularizer
    """
    if norm_reg:
        norm = flux**2
    else:
        norm = 1

    out = -0.5*np.sum((imvec - priorvec) ** 2)
    return out/norm


def spatchgrad(imvec, priorvec, flux, norm_reg=NORM_REGULARIZER):
    """Patch prior gradient
    """
    if norm_reg:
        norm = flux**2
    else:
        norm = 1

    out = -(imvec - priorvec)
    return out/norm

# TODO FIGURE OUT NORMALIZATIONS FOR COMPACT 1 & 2 REGULARIZERS


def scompact(imvec, nx, ny, psize, flux, norm_reg=NORM_REGULARIZER, beam_size=None):
    """I r^2 source size regularizer
    """

    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = flux * (beam_size**2)
    else:
        norm = 1

    im = imvec.reshape(ny, nx)

    xx, yy = np.meshgrid(range(nx), range(ny))
    xx = xx - (nx-1)/2.0
    yy = yy - (ny-1)/2.0
    xxpsize = xx * psize
    yypsize = yy * psize

    x0 = np.sum(np.sum(im * xxpsize))/flux
    y0 = np.sum(np.sum(im * yypsize))/flux

    out = -np.sum(np.sum(im * ((xxpsize - x0)**2 + (yypsize - y0)**2)))
    return out/norm


def scompactgrad(imvec, nx, ny, psize, flux, norm_reg=NORM_REGULARIZER, beam_size=None):
    """Gradient for I r^2 source size regularizer
    """

    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = flux * beam_size**2
    else:
        norm = 1

    im = imvec.reshape(ny, nx)

    xx, yy = np.meshgrid(range(nx), range(ny))
    xx = xx - (nx-1)/2.0
    yy = yy - (ny-1)/2.0
    xxpsize = xx * psize
    yypsize = yy * psize

    x0 = np.sum(np.sum(im * xxpsize))/flux
    y0 = np.sum(np.sum(im * yypsize))/flux

    term1 = np.sum(np.sum(im * ((xxpsize - x0))))
    term2 = np.sum(np.sum(im * ((yypsize - y0))))

    grad = -2*xxpsize*term1 - 2*yypsize*term2 + (xxpsize - x0)**2 + (yypsize - y0)**2

    return -grad.reshape(-1)/norm


def scompact2(imvec, nx, ny, psize, flux, norm_reg=NORM_REGULARIZER, beam_size=None):
    """I^2r^2 source size regularizer
    """

    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = flux**2 * beam_size**2
    else:
        norm = 1

    im = imvec.reshape(ny, nx)

    xx, yy = np.meshgrid(range(nx), range(ny))
    xx = xx - (nx-1)/2.0
    yy = yy - (ny-1)/2.0
    xxpsize = xx * psize
    yypsize = yy * psize

    out = -np.sum(np.sum(im**2 * (xxpsize**2 + yypsize**2)))
    return out/norm


def scompact2grad(imvec, nx, ny, psize, flux, norm_reg=NORM_REGULARIZER, beam_size=None):
    """Gradient for I^2r^2 source size regularizer
    """

    if beam_size is None:
        beam_size = psize
    if norm_reg:
        norm = flux**2 * beam_size**2
    else:
        norm = 1

    im = imvec.reshape(ny, nx)

    xx, yy = np.meshgrid(range(nx), range(ny))
    xx = xx - (nx-1)/2.0
    yy = yy - (ny-1)/2.0
    xxpsize = xx * psize
    yypsize = yy * psize

    grad = -2*im*(xxpsize**2 + yypsize**2)

    return grad.reshape(-1)/norm


def sgauss(imvec, xdim, ydim, psize, major, minor, PA):
    """Gaussian source size regularizer
    """

    # major, minor and PA are all in radians
    phi = PA

    # eigenvalues of covariance matrix
    lambda1 = minor**2./(8.*np.log(2.))
    lambda2 = major**2./(8.*np.log(2.))

    # now compute covariance matrix elements from user inputs
    sigxx_prime = lambda1*(np.cos(phi)**2.) + lambda2*(np.sin(phi)**2.)
    sigyy_prime = lambda1*(np.sin(phi)**2.) + lambda2*(np.cos(phi)**2.)
    sigxy_prime = (lambda2 - lambda1)*np.cos(phi)*np.sin(phi)

    # we get the dimensions and image vector
    im = imvec.reshape(xdim, ydim)
    xlist, ylist = np.meshgrid(range(xdim), range(ydim))
    xlist = xlist - (xdim-1)/2.0
    ylist = ylist - (ydim-1)/2.0

    xx = xlist * psize
    yy = ylist * psize

    # the centroid parameters
    x0 = np.sum(xx*im) / np.sum(im)
    y0 = np.sum(yy*im) / np.sum(im)

    # we calculate the elements of the covariance matrix
    sigxx = (np.sum((xx - x0)**2.*im)/np.sum(im))
    sigyy = (np.sum((yy - y0)**2.*im)/np.sum(im))
    sigxy = (np.sum((xx - x0)*(yy - y0)*im)/np.sum(im))

    # We calculate the regularizer #this line was CHANGED
    rgauss = -((sigxx - sigxx_prime)**2. + (sigyy - sigyy_prime)**2. + 2*(sigxy - sigxy_prime)**2.)
    # normalization will need to be redone, right now requires alpha~1000
    rgauss = rgauss/(major**2. * minor**2.)
    return rgauss


def sgauss_grad(imvec, xdim, ydim, psize, major, minor, PA):
    """Gradient for Gaussian source size regularizer
    """

    # major, minor and PA are all in radians
    phi = PA

    # computing eigenvalues of the covariance matrix
    lambda1 = (minor**2.)/(8.*np.log(2.))
    lambda2 = (major**2.)/(8.*np.log(2.))

    # now compute covariance matrix elements from user inputs

    sigxx_prime = lambda1*(np.cos(phi)**2.) + lambda2*(np.sin(phi)**2.)
    sigyy_prime = lambda1*(np.sin(phi)**2.) + lambda2*(np.cos(phi)**2.)
    sigxy_prime = (lambda2 - lambda1)*np.cos(phi)*np.sin(phi)

    # we get the dimensions and image vector
    im = imvec.reshape(xdim, ydim)
    xlist, ylist = np.meshgrid(range(xdim), range(ydim))
    xlist = xlist - (xdim-1)/2.0
    ylist = ylist - (ydim-1)/2.0

    xx = xlist * psize
    yy = ylist * psize

    # the centroid parameters
    x0 = np.sum(xx*im) / np.sum(im)
    y0 = np.sum(yy*im) / np.sum(im)

    # we calculate the elements of the covariance matrix of the image
    sigxx = (np.sum((xx - x0)**2.*im)/np.sum(im))
    sigyy = (np.sum((yy - y0)**2.*im)/np.sum(im))
    sigxy = (np.sum((xx - x0)*(yy - y0)*im)/np.sum(im))

    # now we compute the gradients of all quantities
    # gradient of centroid
    dx0 = (xx - x0) / np.sum(im)
    dy0 = (yy - y0) / np.sum(im)

    # gradients of covariance matrix elements
    dxx = (((xx - x0)**2. - 2.*(xx - x0)*dx0*im) - sigxx) / np.sum(im)

    dyy = (((yy - y0)**2. - 2.*(yy - y0)*dx0*im) - sigyy) / np.sum(im)

    dxy = (((xx - x0)*(yy - y0) - (yy - y0)*dx0*im - (xx - x0)*dy0*im) - sigxy) / np.sum(im)

    # gradient of the regularizer #this line was CHANGED
    drgauss = (2.*(sigxx - sigxx_prime)*dxx +
               2.*(sigyy - sigyy_prime)*dyy +
               4.*(sigxy - sigxy_prime)*dxy)

    # normalization will need to be redone, right now requires alpha~1000
    drgauss = drgauss/(major**2. * minor**2.)

    return -drgauss.reshape(-1)


##################################################################################################
# Chi^2 Data functions
##################################################################################################
def apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol):
    """apply systematic noise to VISIBILITIES or AMPLITUDES
       data_arr should have fields 't1','t2','u','v','vis','amp','sigma'

       returns: (uv, vis, amp, sigma)
    """

    vtype = ehc.vis_poldict[pol]
    atype = ehc.amp_poldict[pol]
    etype = ehc.sig_poldict[pol]

    t1 = data_arr['t1']
    t2 = data_arr['t2']

    sigma = data_arr[etype]
    amp = data_arr[atype]
    try:
        vis = data_arr[vtype]
    except ValueError:
        vis = amp.astype('c16')

    snrmask = np.abs(amp/sigma) >= snrcut

    if type(systematic_noise) is dict:
        sys_level = np.zeros(len(t1))
        for i in range(len(t1)):
            if t1[i] in systematic_noise.keys():
                t1sys = systematic_noise[t1[i]]
            else:
                t1sys = 0.
            if t2[i] in systematic_noise.keys():
                t2sys = systematic_noise[t2[i]]
            else:
                t2sys = 0.

            if t1sys < 0 or t2sys < 0:
                sys_level[i] = -1
            else:
                sys_level[i] = np.sqrt(t1sys**2 + t2sys**2)
    else:
        sys_level = np.sqrt(2)*systematic_noise*np.ones(len(t1))

    mask = sys_level >= 0.
    mask = snrmask * mask

    sigma = np.linalg.norm([sigma, sys_level*np.abs(amp)], axis=0)[mask]
    vis = vis[mask]
    amp = amp[mask]
    uv = np.hstack((data_arr['u'].reshape(-1, 1), data_arr['v'].reshape(-1, 1)))[mask]
    return (uv, vis, amp, sigma)


def chisqdata_vis(Obsdata, Prior, mask, pol='I', **kwargs):
    """Return the data, sigmas, and fourier matrix for visibilities
    """

    # unpack keyword args
    systematic_noise = kwargs.get('systematic_noise', 0.)
    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')

    # unpack data
    vtype = ehc.vis_poldict[pol]
    atype = ehc.amp_poldict[pol]
    etype = ehc.sig_poldict[pol]
    data_arr = Obsdata.unpack(['t1', 't2', 'u', 'v', vtype, atype, etype], debias=debias)
    (uv, vis, amp, sigma) = apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # make fourier matrix
    A = obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv, pulse=Prior.pulse, mask=mask)

    return (vis, sigma, A)


def chisqdata_amp(Obsdata, Prior, mask, pol='I', **kwargs):
    """Return the data, sigmas, and fourier matrix for visibility amplitudes
    """

    # unpack keyword args
    systematic_noise = kwargs.get('systematic_noise', 0.)
    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')

    # unpack data
    vtype = ehc.vis_poldict[pol]
    atype = ehc.amp_poldict[pol]
    etype = ehc.sig_poldict[pol]
    if (Obsdata.amp is None) or (len(Obsdata.amp) == 0) or pol != 'I':
        data_arr = Obsdata.unpack(['t1', 't2', 'u', 'v', vtype, atype, etype], debias=debias)

    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed amplitude table in amplitude chi^2!")
        if not type(Obsdata.amp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed amplitude table is not a numpy rec array!")
        data_arr = Obsdata.amp

    # apply systematic noise and SNR cut
    # TODO -- after pre-computed??
    (uv, vis, amp, sigma) = apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # make fourier matrix
    A = obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv, pulse=Prior.pulse, mask=mask)

    return (amp, sigma, A)


def chisqdata_bs(Obsdata, Prior, mask, pol='I', **kwargs):
    """return the data, sigmas, and fourier matrices for bispectra
    """

    # unpack keyword args
    # systematic_noise = kwargs.get('systematic_noise',0.)
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    weighting = kwargs.get('weighting', 'natural')

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.bispec is None) or (len(Obsdata.bispec) == 0) or pol != 'I':
        biarr = Obsdata.bispectra(mode="all", vtype=vtype, count=count, snrcut=snrcut)

    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed bispectrum table in cphase chi^2!")
        if not type(Obsdata.bispec) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed bispectrum table is not a numpy rec array!")
        biarr = Obsdata.bispec
        # reduce to a minimal set
        if count != 'max':
            biarr = obsh.reduce_tri_minimal(Obsdata, biarr)

    uv1 = np.hstack((biarr['u1'].reshape(-1, 1), biarr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((biarr['u2'].reshape(-1, 1), biarr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((biarr['u3'].reshape(-1, 1), biarr['v3'].reshape(-1, 1)))
    bi = biarr['bispec']
    sigma = biarr['sigmab']

    # add systematic noise
    # sigma = np.linalg.norm([biarr['sigmab'], systematic_noise*np.abs(biarr['bispec'])], axis=0)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # make fourier matrices
    A3 = (obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv1, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv2, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv3, pulse=Prior.pulse, mask=mask)
         )

    return (bi, sigma, A3)


def chisqdata_cphase(Obsdata, Prior, mask, pol='I', **kwargs):
    """Return the data, sigmas, and fourier matrices for closure phases
    """

    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    uv_min = kwargs.get('cp_uv_min', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    systematic_cphase_noise = kwargs.get('systematic_cphase_noise', 0.)
    weighting = kwargs.get('weighting', 'natural')

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.cphase is None) or (len(Obsdata.cphase) == 0) or pol != 'I':
        clphasearr = Obsdata.c_phases(mode="all", vtype=vtype,
                                      count=count, uv_min=uv_min, snrcut=snrcut)
    else:  # TODO precomputed with not Stokes I
        print("Using pre-computed cphase table in cphase chi^2!")
        if not type(Obsdata.cphase) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed closure phase table is not a numpy rec array!")
        clphasearr = Obsdata.cphase
        # reduce to a minimal set
        if count != 'max':
            clphasearr = obsh.reduce_tri_minimal(Obsdata, clphasearr)

    uv1 = np.hstack((clphasearr['u1'].reshape(-1, 1), clphasearr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clphasearr['u2'].reshape(-1, 1), clphasearr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clphasearr['u3'].reshape(-1, 1), clphasearr['v3'].reshape(-1, 1)))
    clphase = clphasearr['cphase']
    sigma = clphasearr['sigmacp']

    # add systematic cphase noise (in DEGREES)
    sigma = np.linalg.norm([sigma, systematic_cphase_noise*np.ones(len(sigma))], axis=0)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # make fourier matrices
    A3 = (obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv1, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv2, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv3, pulse=Prior.pulse, mask=mask)
          )
    return (clphase, sigma, A3)


def chisqdata_cphase_diag(Obsdata, Prior, mask, pol='I', **kwargs):
    """Return the data, sigmas, and fourier matrices for diagonalized closure phases
    """

    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    uv_min = kwargs.get('cp_uv_min', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    clphasearr = Obsdata.c_phases_diag(vtype=vtype, count=count, snrcut=snrcut, uv_min=uv_min)

    # loop over timestamps
    clphase_diag = []
    sigma_diag = []
    A3_diag = []
    tform_mats = []
    for ic, cl in enumerate(clphasearr):

        # get diagonalized closure phases and errors
        clphase_diag.append(cl[0]['cphase'])
        sigma_diag.append(cl[0]['sigmacp'])

        # get uv arrays
        u1 = cl[2][:, 0].astype('float')
        v1 = cl[3][:, 0].astype('float')
        uv1 = np.hstack((u1.reshape(-1, 1), v1.reshape(-1, 1)))

        u2 = cl[2][:, 1].astype('float')
        v2 = cl[3][:, 1].astype('float')
        uv2 = np.hstack((u2.reshape(-1, 1), v2.reshape(-1, 1)))

        u3 = cl[2][:, 2].astype('float')
        v3 = cl[3][:, 2].astype('float')
        uv3 = np.hstack((u3.reshape(-1, 1), v3.reshape(-1, 1)))

        # compute Fourier matrices
        A3 = (obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv1, pulse=Prior.pulse, mask=mask),
              obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv2, pulse=Prior.pulse, mask=mask),
              obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv3, pulse=Prior.pulse, mask=mask)
              )
        A3_diag.append(A3)

        # get transformation matrix for this timestamp
        tform_mats.append(cl[4].astype('float'))

    # combine Fourier and transformation matrices into tuple for outputting
    Amatrices = (np.array(A3_diag), np.array(tform_mats))

    return (np.array(clphase_diag), np.array(sigma_diag), Amatrices)


def chisqdata_camp(Obsdata, Prior, mask, pol='I', **kwargs):
    """Return the data, sigmas, and fourier matrices for closure amplitudes
    """
    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')

    # unpack data & mask low snr points
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.camp is None) or (len(Obsdata.camp) == 0) or pol != 'I':
        clamparr = Obsdata.c_amplitudes(mode='all', count=count,
                                        vtype=vtype, ctype='camp', debias=debias, snrcut=snrcut)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed closure amplitude table in closure amplitude chi^2!")
        if not type(Obsdata.camp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed closure amplitude table is not a numpy rec array!")
        clamparr = Obsdata.camp
        # reduce to a minimal set
        if count != 'max':
            clamparr = obsh.reduce_quad_minimal(Obsdata, clamparr, ctype='camp')

    uv1 = np.hstack((clamparr['u1'].reshape(-1, 1), clamparr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clamparr['u2'].reshape(-1, 1), clamparr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clamparr['u3'].reshape(-1, 1), clamparr['v3'].reshape(-1, 1)))
    uv4 = np.hstack((clamparr['u4'].reshape(-1, 1), clamparr['v4'].reshape(-1, 1)))
    clamp = clamparr['camp']
    sigma = clamparr['sigmaca']

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # make fourier matrices
    A4 = (obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv1, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv2, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv3, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv4, pulse=Prior.pulse, mask=mask)
          )

    return (clamp, sigma, A4)


def chisqdata_logcamp(Obsdata, Prior, mask, pol='I', **kwargs):
    """Return the data, sigmas, and fourier matrices for log closure amplitudes
    """
    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')

    # unpack data & mask low snr points
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.logcamp is None) or (len(Obsdata.logcamp) == 0) or pol != 'I':
        clamparr = Obsdata.c_amplitudes(mode='all', count=count,
                                        vtype=vtype, ctype='logcamp', debias=debias, snrcut=snrcut)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed log closure amplitude table in log closure amplitude chi^2!")
        if not type(Obsdata.logcamp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed log closure amplitude table is not a numpy rec array!")
        clamparr = Obsdata.logcamp
        # reduce to a minimal set
        if count != 'max':
            clamparr = obsh.reduce_quad_minimal(Obsdata, clamparr, ctype='logcamp')

    uv1 = np.hstack((clamparr['u1'].reshape(-1, 1), clamparr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clamparr['u2'].reshape(-1, 1), clamparr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clamparr['u3'].reshape(-1, 1), clamparr['v3'].reshape(-1, 1)))
    uv4 = np.hstack((clamparr['u4'].reshape(-1, 1), clamparr['v4'].reshape(-1, 1)))
    clamp = clamparr['camp']
    sigma = clamparr['sigmaca']

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # make fourier matrices
    A4 = (obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv1, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv2, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv3, pulse=Prior.pulse, mask=mask),
          obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv4, pulse=Prior.pulse, mask=mask)
          )

    return (clamp, sigma, A4)


def chisqdata_logcamp_diag(Obsdata, Prior, mask, pol='I', **kwargs):
    """Return the data, sigmas, and fourier matrices for diagonalized log closure amplitudes
    """
    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)

    # unpack data & mask low snr points
    vtype = ehc.vis_poldict[pol]
    clamparr = Obsdata.c_log_amplitudes_diag(vtype=vtype, count=count, debias=debias, snrcut=snrcut)

    # loop over timestamps
    clamp_diag = []
    sigma_diag = []
    A4_diag = []
    tform_mats = []
    for ic, cl in enumerate(clamparr):

        # get diagonalized log closure amplitudes and errors
        clamp_diag.append(cl[0]['camp'])
        sigma_diag.append(cl[0]['sigmaca'])

        # get uv arrays
        u1 = cl[2][:, 0].astype('float')
        v1 = cl[3][:, 0].astype('float')
        uv1 = np.hstack((u1.reshape(-1, 1), v1.reshape(-1, 1)))

        u2 = cl[2][:, 1].astype('float')
        v2 = cl[3][:, 1].astype('float')
        uv2 = np.hstack((u2.reshape(-1, 1), v2.reshape(-1, 1)))

        u3 = cl[2][:, 2].astype('float')
        v3 = cl[3][:, 2].astype('float')
        uv3 = np.hstack((u3.reshape(-1, 1), v3.reshape(-1, 1)))

        u4 = cl[2][:, 3].astype('float')
        v4 = cl[3][:, 3].astype('float')
        uv4 = np.hstack((u4.reshape(-1, 1), v4.reshape(-1, 1)))

        # compute Fourier matrices
        A4 = (obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv1, pulse=Prior.pulse, mask=mask),
              obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv2, pulse=Prior.pulse, mask=mask),
              obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv3, pulse=Prior.pulse, mask=mask),
              obsh.ftmatrix(Prior.psize, Prior.xdim, Prior.ydim, uv4, pulse=Prior.pulse, mask=mask)
              )
        A4_diag.append(A4)

        # get transformation matrix for this timestamp
        tform_mats.append(cl[4].astype('float'))

    # combine Fourier and transformation matrices into tuple for outputting
    Amatrices = (np.array(A4_diag), np.array(tform_mats))

    return (np.array(clamp_diag), np.array(sigma_diag), Amatrices)

##################################################################################################
# FFT Chi^2 Data functions
##################################################################################################


def chisqdata_vis_fft(Obsdata, Prior, pol='I', **kwargs):
    """Return the data, sigmas, uv points, and FFT info for visibilities
    """

    # unpack keyword args
    systematic_noise = kwargs.get('systematic_noise', 0.)
    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    conv_func = kwargs.get('conv_func', ehc.GRIDDER_CONV_FUNC_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)
    order = kwargs.get('order', ehc.FFT_INTERP_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    atype = ehc.amp_poldict[pol]
    etype = ehc.sig_poldict[pol]
    data_arr = Obsdata.unpack(['t1', 't2', 'u', 'v', vtype, atype, etype], debias=debias)
    (uv, vis, amp, sigma) = apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # prepare image and fft info objects
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    im_info = obsh.ImInfo(Prior.xdim, Prior.ydim, npad, Prior.psize, Prior.pulse)
    gs_info = obsh.make_gridder_and_sampler_info(
        im_info, uv, conv_func=conv_func, p_rad=p_rad, order=order)

    sampler_info_list = [gs_info[0]]
    gridder_info_list = [gs_info[1]]
    A = (im_info, sampler_info_list, gridder_info_list)

    return (vis, sigma, A)


def chisqdata_amp_fft(Obsdata, Prior, pol='I', **kwargs):
    """Return the data, sigmas, uv points, and FFT info for visibility amplitudes
    """

    # unpack keyword args
    systematic_noise = kwargs.get('systematic_noise', 0.)
    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    conv_func = kwargs.get('conv_func', ehc.GRIDDER_CONV_FUNC_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)
    order = kwargs.get('order', ehc.FFT_INTERP_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    atype = ehc.amp_poldict[pol]
    etype = ehc.sig_poldict[pol]
    if (Obsdata.amp is None) or (len(Obsdata.amp) == 0) or pol != 'I':
        data_arr = Obsdata.unpack(['t1', 't2', 'u', 'v', vtype, atype, etype], debias=debias)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed amplitude table in amplitude chi^2!")
        if not type(Obsdata.amp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed amplitude table is not a numpy rec array!")
        data_arr = Obsdata.amp

    # apply systematic noise
    (uv, vis, amp, sigma) = apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # prepare image and fft info objects
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    im_info = obsh.ImInfo(Prior.xdim, Prior.ydim, npad, Prior.psize, Prior.pulse)
    gs_info = obsh.make_gridder_and_sampler_info(im_info, uv,
                                                 conv_func=conv_func, p_rad=p_rad, order=order)

    sampler_info_list = [gs_info[0]]
    gridder_info_list = [gs_info[1]]
    A = (im_info, sampler_info_list, gridder_info_list)

    return (amp, sigma, A)


def chisqdata_bs_fft(Obsdata, Prior, pol='I', **kwargs):
    """Return the data, sigmas, uv points, and FFT info for bispectra
    """

    # unpack keyword args
    # systematic_noise = kwargs.get('systematic_noise',0.)
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    conv_func = kwargs.get('conv_func', ehc.GRIDDER_CONV_FUNC_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)
    order = kwargs.get('order', ehc.FFT_INTERP_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.bispec is None) or (len(Obsdata.bispec) == 0) or pol != 'I':
        biarr = Obsdata.bispectra(mode="all", vtype=vtype, count=count, snrcut=snrcut)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed bispectrum table in cphase chi^2!")
        if not type(Obsdata.bispec) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed bispectrum table is not a numpy rec array!")
        biarr = Obsdata.bispec
        # reduce to a minimal set
        if count != 'max':
            biarr = obsh.reduce_tri_minimal(Obsdata, biarr)

    uv1 = np.hstack((biarr['u1'].reshape(-1, 1), biarr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((biarr['u2'].reshape(-1, 1), biarr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((biarr['u3'].reshape(-1, 1), biarr['v3'].reshape(-1, 1)))
    bi = biarr['bispec']
    sigma = biarr['sigmab']

    # add systematic noise
    # sigma = np.linalg.norm([biarr['sigmab'], systematic_noise*np.abs(biarr['bispec'])], axis=0)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # prepare image and fft info objects
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    im_info = obsh.ImInfo(Prior.xdim, Prior.ydim, npad, Prior.psize, Prior.pulse)
    gs_info1 = obsh.make_gridder_and_sampler_info(im_info, uv1,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info2 = obsh.make_gridder_and_sampler_info(im_info, uv2,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info3 = obsh.make_gridder_and_sampler_info(im_info, uv3,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)

    sampler_info_list = [gs_info1[0], gs_info2[0], gs_info3[0]]
    gridder_info_list = [gs_info1[1], gs_info2[1], gs_info3[1]]
    A = (im_info, sampler_info_list, gridder_info_list)

    return (bi, sigma, A)


def chisqdata_cphase_fft(Obsdata, Prior, pol='I', **kwargs):
    """Return the data, sigmas, uv points, and FFT info for closure phases
    """
    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    uv_min = kwargs.get('cp_uv_min', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    weighting = kwargs.get('weighting', 'natural')
    systematic_cphase_noise = kwargs.get('systematic_cphase_noise', 0.)
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    conv_func = kwargs.get('conv_func', ehc.GRIDDER_CONV_FUNC_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)
    order = kwargs.get('order', ehc.FFT_INTERP_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.cphase is None) or (len(Obsdata.cphase) == 0) or pol != 'I':
        clphasearr = Obsdata.c_phases(mode="all", vtype=vtype,
                                      count=count, uv_min=uv_min, snrcut=snrcut)
    else:  # TODO precomputed with not Stokes I
        print("Using pre-computed cphase table in cphase chi^2!")
        if not type(Obsdata.cphase) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed closure phase table is not a numpy rec array!")
        clphasearr = Obsdata.cphase
        # reduce to a minimal set
        if count != 'max':
            clphasearr = obsh.reduce_tri_minimal(Obsdata, clphasearr)

    uv1 = np.hstack((clphasearr['u1'].reshape(-1, 1), clphasearr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clphasearr['u2'].reshape(-1, 1), clphasearr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clphasearr['u3'].reshape(-1, 1), clphasearr['v3'].reshape(-1, 1)))
    clphase = clphasearr['cphase']
    sigma = clphasearr['sigmacp']

    # add systematic cphase noise (in DEGREES)
    sigma = np.linalg.norm([sigma, systematic_cphase_noise*np.ones(len(sigma))], axis=0)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # prepare image and fft info objects
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    im_info = obsh.ImInfo(Prior.xdim, Prior.ydim, npad, Prior.psize, Prior.pulse)
    gs_info1 = obsh.make_gridder_and_sampler_info(im_info, uv1,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info2 = obsh.make_gridder_and_sampler_info(im_info, uv2,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info3 = obsh.make_gridder_and_sampler_info(im_info, uv3,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)

    sampler_info_list = [gs_info1[0], gs_info2[0], gs_info3[0]]
    gridder_info_list = [gs_info1[1], gs_info2[1], gs_info3[1]]
    A = (im_info, sampler_info_list, gridder_info_list)

    return (clphase, sigma, A)


def chisqdata_cphase_diag_fft(Obsdata, Prior, pol='I', **kwargs):
    """Return the data, sigmas, uv points, and FFT info for diagonalized closure phases
    """
    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    uv_min = kwargs.get('cp_uv_min', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    conv_func = kwargs.get('conv_func', ehc.GRIDDER_CONV_FUNC_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)
    order = kwargs.get('order', ehc.FFT_INTERP_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    clphasearr = Obsdata.c_phases_diag(vtype=vtype, count=count, snrcut=snrcut, uv_min=uv_min)

    # loop over timestamps
    clphase_diag = []
    sigma_diag = []
    tform_mats = []
    u1 = []
    v1 = []
    u2 = []
    v2 = []
    u3 = []
    v3 = []
    for ic, cl in enumerate(clphasearr):

        # get diagonalized closure phases and errors
        clphase_diag.append(cl[0]['cphase'])
        sigma_diag.append(cl[0]['sigmacp'])

        # get u and v values
        u1.append(cl[2][:, 0].astype('float'))
        v1.append(cl[3][:, 0].astype('float'))
        u2.append(cl[2][:, 1].astype('float'))
        v2.append(cl[3][:, 1].astype('float'))
        u3.append(cl[2][:, 2].astype('float'))
        v3.append(cl[3][:, 2].astype('float'))

        # get transformation matrix for this timestamp
        tform_mats.append(cl[4].astype('float'))

    # fix formatting of arrays
    u1 = np.concatenate(u1)
    v1 = np.concatenate(v1)
    u2 = np.concatenate(u2)
    v2 = np.concatenate(v2)
    u3 = np.concatenate(u3)
    v3 = np.concatenate(v3)

    # get uv arrays
    uv1 = np.hstack((u1.reshape(-1, 1), v1.reshape(-1, 1)))
    uv2 = np.hstack((u2.reshape(-1, 1), v2.reshape(-1, 1)))
    uv3 = np.hstack((u3.reshape(-1, 1), v3.reshape(-1, 1)))

    # prepare image and fft info objects
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    im_info = obsh.ImInfo(Prior.xdim, Prior.ydim, npad, Prior.psize, Prior.pulse)
    gs_info1 = obsh.make_gridder_and_sampler_info(im_info, uv1,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info2 = obsh.make_gridder_and_sampler_info(im_info, uv2,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info3 = obsh.make_gridder_and_sampler_info(im_info, uv3,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)

    sampler_info_list = [gs_info1[0], gs_info2[0], gs_info3[0]]
    gridder_info_list = [gs_info1[1], gs_info2[1], gs_info3[1]]
    A3 = (im_info, sampler_info_list, gridder_info_list)

    # combine Fourier and transformation matrices into tuple for outputting
    Amatrices = (A3, np.array(tform_mats))

    return (np.array(clphase_diag), np.array(sigma_diag), Amatrices)


def chisqdata_camp_fft(Obsdata, Prior, pol='I', **kwargs):
    """Return the data, sigmas, uv points, and FFT info for closure amplitudes
    """

    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    conv_func = kwargs.get('conv_func', ehc.GRIDDER_CONV_FUNC_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)
    order = kwargs.get('order', ehc.FFT_INTERP_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.camp is None) or (len(Obsdata.camp) == 0) or pol != 'I':
        clamparr = Obsdata.c_amplitudes(mode='all', count=count,
                                        vtype=vtype, ctype='camp', debias=debias, snrcut=snrcut)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed closure amplitude table in closure amplitude chi^2!")
        if not type(Obsdata.camp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed closure amplitude table is not a numpy rec array!")
        clamparr = Obsdata.camp
        # reduce to a minimal set
        if count != 'max':
            clamparr = obsh.reduce_quad_minimal(Obsdata, clamparr, ctype='camp')

    uv1 = np.hstack((clamparr['u1'].reshape(-1, 1), clamparr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clamparr['u2'].reshape(-1, 1), clamparr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clamparr['u3'].reshape(-1, 1), clamparr['v3'].reshape(-1, 1)))
    uv4 = np.hstack((clamparr['u4'].reshape(-1, 1), clamparr['v4'].reshape(-1, 1)))
    clamp = clamparr['camp']
    sigma = clamparr['sigmaca']

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # prepare image and fft info objects
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    im_info = obsh.ImInfo(Prior.xdim, Prior.ydim, npad, Prior.psize, Prior.pulse)
    gs_info1 = obsh.make_gridder_and_sampler_info(im_info, uv1,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info2 = obsh.make_gridder_and_sampler_info(im_info, uv2,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info3 = obsh.make_gridder_and_sampler_info(im_info, uv3,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info4 = obsh.make_gridder_and_sampler_info(im_info, uv4,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)

    sampler_info_list = [gs_info1[0], gs_info2[0], gs_info3[0], gs_info4[0]]
    gridder_info_list = [gs_info1[1], gs_info2[1], gs_info3[1], gs_info4[1]]
    A = (im_info, sampler_info_list, gridder_info_list)

    return (clamp, sigma, A)


def chisqdata_logcamp_fft(Obsdata, Prior, pol='I', **kwargs):
    """Return the data, sigmas, uv points, and FFT info for log closure amplitudes
    """
    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    conv_func = kwargs.get('conv_func', ehc.GRIDDER_CONV_FUNC_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)
    order = kwargs.get('order', ehc.FFT_INTERP_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.logcamp is None) or (len(Obsdata.logcamp) == 0) or pol != 'I':
        clamparr = Obsdata.c_amplitudes(mode='all', count=count,
                                        vtype=vtype, ctype='logcamp', debias=debias, snrcut=snrcut)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed log closure amplitude table in log closure amplitude chi^2!")
        if not type(Obsdata.logcamp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed log closure amplitude table is not a numpy rec array!")
        clamparr = Obsdata.logcamp
        # reduce to a minimal set
        if count != 'max':
            clamparr = obsh.reduce_quad_minimal(Obsdata, clamparr, ctype='logcamp')

    uv1 = np.hstack((clamparr['u1'].reshape(-1, 1), clamparr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clamparr['u2'].reshape(-1, 1), clamparr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clamparr['u3'].reshape(-1, 1), clamparr['v3'].reshape(-1, 1)))
    uv4 = np.hstack((clamparr['u4'].reshape(-1, 1), clamparr['v4'].reshape(-1, 1)))
    clamp = clamparr['camp']
    sigma = clamparr['sigmaca']

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # prepare image and fft info objects
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    im_info = obsh.ImInfo(Prior.xdim, Prior.ydim, npad, Prior.psize, Prior.pulse)
    gs_info1 = obsh.make_gridder_and_sampler_info(im_info, uv1,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info2 = obsh.make_gridder_and_sampler_info(im_info, uv2,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info3 = obsh.make_gridder_and_sampler_info(im_info, uv3,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info4 = obsh.make_gridder_and_sampler_info(im_info, uv4,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)

    sampler_info_list = [gs_info1[0], gs_info2[0], gs_info3[0], gs_info4[0]]
    gridder_info_list = [gs_info1[1], gs_info2[1], gs_info3[1], gs_info4[1]]
    A = (im_info, sampler_info_list, gridder_info_list)

    return (clamp, sigma, A)


def chisqdata_logcamp_diag_fft(Obsdata, Prior, pol='I', **kwargs):
    """Return the data, sigmas, uv points, and FFT info for diagonalized log closure amplitudes
    """
    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    conv_func = kwargs.get('conv_func', ehc.GRIDDER_CONV_FUNC_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)
    order = kwargs.get('order', ehc.FFT_INTERP_DEFAULT)

    # unpack data & mask low snr points
    vtype = ehc.vis_poldict[pol]
    clamparr = Obsdata.c_log_amplitudes_diag(vtype=vtype, count=count,
                                             debias=debias, snrcut=snrcut)

    # loop over timestamps
    clamp_diag = []
    sigma_diag = []
    tform_mats = []
    u1 = []
    v1 = []
    u2 = []
    v2 = []
    u3 = []
    v3 = []
    u4 = []
    v4 = []
    for ic, cl in enumerate(clamparr):

        # get diagonalized log closure amplitudes and errors
        clamp_diag.append(cl[0]['camp'])
        sigma_diag.append(cl[0]['sigmaca'])

        # get u and v values
        u1.append(cl[2][:, 0].astype('float'))
        v1.append(cl[3][:, 0].astype('float'))
        u2.append(cl[2][:, 1].astype('float'))
        v2.append(cl[3][:, 1].astype('float'))
        u3.append(cl[2][:, 2].astype('float'))
        v3.append(cl[3][:, 2].astype('float'))
        u4.append(cl[2][:, 3].astype('float'))
        v4.append(cl[3][:, 3].astype('float'))

        # get transformation matrix for this timestamp
        tform_mats.append(cl[4].astype('float'))

    # fix formatting of arrays
    u1 = np.concatenate(u1)
    v1 = np.concatenate(v1)
    u2 = np.concatenate(u2)
    v2 = np.concatenate(v2)
    u3 = np.concatenate(u3)
    v3 = np.concatenate(v3)
    u4 = np.concatenate(u4)
    v4 = np.concatenate(v4)

    # get uv arrays
    uv1 = np.hstack((u1.reshape(-1, 1), v1.reshape(-1, 1)))
    uv2 = np.hstack((u2.reshape(-1, 1), v2.reshape(-1, 1)))
    uv3 = np.hstack((u3.reshape(-1, 1), v3.reshape(-1, 1)))
    uv4 = np.hstack((u4.reshape(-1, 1), v4.reshape(-1, 1)))

    # prepare image and fft info objects
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    im_info = obsh.ImInfo(Prior.xdim, Prior.ydim, npad, Prior.psize, Prior.pulse)
    gs_info1 = obsh.make_gridder_and_sampler_info(im_info, uv1,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info2 = obsh.make_gridder_and_sampler_info(im_info, uv2,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info3 = obsh.make_gridder_and_sampler_info(im_info, uv3,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)
    gs_info4 = obsh.make_gridder_and_sampler_info(im_info, uv4,
                                                  conv_func=conv_func, p_rad=p_rad, order=order)

    sampler_info_list = [gs_info1[0], gs_info2[0], gs_info3[0], gs_info4[0]]
    gridder_info_list = [gs_info1[1], gs_info2[1], gs_info3[1], gs_info4[1]]
    A = (im_info, sampler_info_list, gridder_info_list)

    # combine Fourier and transformation matrices into tuple for outputting
    Amatrices = (A, np.array(tform_mats))

    return (np.array(clamp_diag), np.array(sigma_diag), Amatrices)

##################################################################################################
# NFFT Chi^2 Data functions
##################################################################################################


def chisqdata_vis_nfft(Obsdata, Prior, pol='I', **kwargs):
    """Return the visibilities, sigmas, uv points, and nfft info
    """
    if (Prior.xdim % 2 or Prior.ydim % 2):
        raise Exception("NFFT doesn't work with odd image dimensions!")

    # unpack keyword args
    systematic_noise = kwargs.get('systematic_noise', 0.)
    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    atype = ehc.amp_poldict[pol]
    etype = ehc.sig_poldict[pol]
    data_arr = Obsdata.unpack(['t1', 't2', 'u', 'v', vtype, atype, etype], debias=debias)
    (uv, vis, amp, sigma) = apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # get NFFT info
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    A1 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv)
    A = [A1]

    return (vis, sigma, A)


def chisqdata_amp_nfft(Obsdata, Prior, pol='I', **kwargs):
    """Return the amplitudes, sigmas, uv points, and nfft info
    """
    if (Prior.xdim % 2 or Prior.ydim % 2):
        raise Exception("NFFT doesn't work with odd image dimensions!")

    # unpack keyword args
    systematic_noise = kwargs.get('systematic_noise', 0.)
    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    atype = ehc.amp_poldict[pol]
    etype = ehc.sig_poldict[pol]
    if (Obsdata.amp is None) or (len(Obsdata.amp) == 0) or pol != 'I':
        data_arr = Obsdata.unpack(['t1', 't2', 'u', 'v', vtype, atype, etype], debias=debias)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed amplitude table in amplitude chi^2!")
        if not type(Obsdata.amp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed amplitude table is not a numpy rec array!")
        data_arr = Obsdata.amp

    # apply systematic noise
    (uv, vis, amp, sigma) = apply_systematic_noise_snrcut(data_arr, systematic_noise, snrcut, pol)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # get NFFT info
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    A1 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv)
    A = [A1]

    return (amp, sigma, A)


def chisqdata_bs_nfft(Obsdata, Prior, pol='I', **kwargs):
    """Return the bispectra, sigmas, uv points, and nfft info
    """
    if (Prior.xdim % 2 or Prior.ydim % 2):
        raise Exception("NFFT doesn't work with odd image dimensions!")

    # unpack keyword args
    # systematic_noise = kwargs.get('systematic_noise',0.) 
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.bispec is None) or (len(Obsdata.bispec) == 0) or pol != 'I':
        biarr = Obsdata.bispectra(mode="all", vtype=vtype, count=count, snrcut=snrcut)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed bispectrum table in cphase chi^2!")
        if not type(Obsdata.bispec) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed bispectrum table is not a numpy rec array!")
        biarr = Obsdata.bispec
        # reduce to a minimal set
        if count != 'max':
            biarr = obsh.reduce_tri_minimal(Obsdata, biarr)

    uv1 = np.hstack((biarr['u1'].reshape(-1, 1), biarr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((biarr['u2'].reshape(-1, 1), biarr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((biarr['u3'].reshape(-1, 1), biarr['v3'].reshape(-1, 1)))
    bi = biarr['bispec']
    sigma = biarr['sigmab']

    # add systematic noise
    # sigma = np.linalg.norm([biarr['sigmab'], systematic_noise*np.abs(biarr['bispec'])], axis=0)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # get NFFT info
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    A1 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv1)
    A2 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv2)
    A3 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv3)
    A = [A1, A2, A3]

    return (bi, sigma, A)


def chisqdata_cphase_nfft(Obsdata, Prior, pol='I', **kwargs):
    """Return the closure phases, sigmas, uv points, and nfft info
    """
    if (Prior.xdim % 2 or Prior.ydim % 2):
        raise Exception("NFFT doesn't work with odd image dimensions!")

    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    uv_min = kwargs.get('cp_uv_min', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    weighting = kwargs.get('weighting', 'natural')
    systematic_cphase_noise = kwargs.get('systematic_cphase_noise', 0.)
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.cphase is None) or (len(Obsdata.cphase) == 0) or pol != 'I':
        clphasearr = Obsdata.c_phases(mode="all", vtype=vtype,
                                      count=count, uv_min=uv_min, snrcut=snrcut)
    else:  # TODO precomputed with not Stokes I
        print("Using pre-computed cphase table in cphase chi^2!")
        if not type(Obsdata.cphase) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed closure phase table is not a numpy rec array!")
        clphasearr = Obsdata.cphase
        # reduce to a minimal set
        if count != 'max':
            clphasearr = obsh.reduce_tri_minimal(Obsdata, clphasearr)

    uv1 = np.hstack((clphasearr['u1'].reshape(-1, 1), clphasearr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clphasearr['u2'].reshape(-1, 1), clphasearr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clphasearr['u3'].reshape(-1, 1), clphasearr['v3'].reshape(-1, 1)))
    clphase = clphasearr['cphase']
    sigma = clphasearr['sigmacp']

    # add systematic cphase noise (in DEGREES)
    sigma = np.linalg.norm([sigma, systematic_cphase_noise*np.ones(len(sigma))], axis=0)

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # get NFFT info
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    A1 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv1)
    A2 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv2)
    A3 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv3)
    A = [A1, A2, A3]

    return (clphase, sigma, A)


def chisqdata_cphase_diag_nfft(Obsdata, Prior, pol='I', **kwargs):
    """Return the diagonalized closure phases, sigmas, uv points, and nfft info
    """
    if (Prior.xdim % 2 or Prior.ydim % 2):
        raise Exception("NFFT doesn't work with odd image dimensions!")

    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    uv_min = kwargs.get('cp_uv_min', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    clphasearr = Obsdata.c_phases_diag(vtype=vtype, count=count, snrcut=snrcut, uv_min=uv_min)

    # loop over timestamps
    clphase_diag = []
    sigma_diag = []
    tform_mats = []
    u1 = []
    v1 = []
    u2 = []
    v2 = []
    u3 = []
    v3 = []
    for ic, cl in enumerate(clphasearr):

        # get diagonalized closure phases and errors
        clphase_diag.append(cl[0]['cphase'])
        sigma_diag.append(cl[0]['sigmacp'])

        # get u and v values
        u1.append(cl[2][:, 0].astype('float'))
        v1.append(cl[3][:, 0].astype('float'))
        u2.append(cl[2][:, 1].astype('float'))
        v2.append(cl[3][:, 1].astype('float'))
        u3.append(cl[2][:, 2].astype('float'))
        v3.append(cl[3][:, 2].astype('float'))

        # get transformation matrix for this timestamp
        tform_mats.append(cl[4].astype('float'))

    # fix formatting of arrays
    u1 = np.concatenate(u1)
    v1 = np.concatenate(v1)
    u2 = np.concatenate(u2)
    v2 = np.concatenate(v2)
    u3 = np.concatenate(u3)
    v3 = np.concatenate(v3)

    # get uv arrays
    uv1 = np.hstack((u1.reshape(-1, 1), v1.reshape(-1, 1)))
    uv2 = np.hstack((u2.reshape(-1, 1), v2.reshape(-1, 1)))
    uv3 = np.hstack((u3.reshape(-1, 1), v3.reshape(-1, 1)))

    # get NFFT info
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    A1 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv1)
    A2 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv2)
    A3 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv3)
    A = [A1, A2, A3]

    # combine Fourier and transformation matrices into tuple for outputting
    Amatrices = (A, np.array(tform_mats))

    return (np.array(clphase_diag), np.array(sigma_diag), Amatrices)


def chisqdata_camp_nfft(Obsdata, Prior, pol='I', **kwargs):
    """Return the closure amplitudes, sigmas, uv points, and nfft info
    """
    if (Prior.xdim % 2 or Prior.ydim % 2):
        raise Exception("NFFT doesn't work with odd image dimensions!")

    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.camp is None) or (len(Obsdata.camp) == 0) or pol != 'I':
        clamparr = Obsdata.c_amplitudes(mode='all', count=count,
                                        vtype=vtype, ctype='camp', debias=debias, snrcut=snrcut)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed closure amplitude table in closure amplitude chi^2!")
        if not type(Obsdata.camp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed closure amplitude table is not a numpy rec array!")
        clamparr = Obsdata.camp
        # reduce to a minimal set
        if count != 'max':
            clamparr = obsh.reduce_quad_minimal(Obsdata, clamparr, ctype='camp')

    uv1 = np.hstack((clamparr['u1'].reshape(-1, 1), clamparr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clamparr['u2'].reshape(-1, 1), clamparr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clamparr['u3'].reshape(-1, 1), clamparr['v3'].reshape(-1, 1)))
    uv4 = np.hstack((clamparr['u4'].reshape(-1, 1), clamparr['v4'].reshape(-1, 1)))
    clamp = clamparr['camp']
    sigma = clamparr['sigmaca']

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # get NFFT info
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    A1 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv1)
    A2 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv2)
    A3 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv3)
    A4 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv4)
    A = [A1, A2, A3, A4]

    return (clamp, sigma, A)


def chisqdata_logcamp_nfft(Obsdata, Prior, pol='I', **kwargs):
    """Return the log closure amplitudes, sigmas, uv points, and nfft info
    """
    if (Prior.xdim % 2 or Prior.ydim % 2):
        raise Exception("NFFT doesn't work with odd image dimensions!")

    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    weighting = kwargs.get('weighting', 'natural')
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)

    # unpack data
    vtype = ehc.vis_poldict[pol]
    if (Obsdata.logcamp is None) or (len(Obsdata.logcamp) == 0) or pol != 'I':
        clamparr = Obsdata.c_amplitudes(mode='all', count=count,
                                        vtype=vtype, ctype='logcamp', debias=debias, snrcut=snrcut)
    else:  # TODO -- pre-computed  with not stokes I?
        print("Using pre-computed log closure amplitude table in log closure amplitude chi^2!")
        if not type(Obsdata.logcamp) in [np.ndarray, np.recarray]:
            raise Exception("pre-computed log closure amplitude table is not a numpy rec array!")
        clamparr = Obsdata.logcamp
        # reduce to a minimal set
        if count != 'max':
            clamparr = obsh.reduce_quad_minimal(Obsdata, clamparr, ctype='logcamp')

    uv1 = np.hstack((clamparr['u1'].reshape(-1, 1), clamparr['v1'].reshape(-1, 1)))
    uv2 = np.hstack((clamparr['u2'].reshape(-1, 1), clamparr['v2'].reshape(-1, 1)))
    uv3 = np.hstack((clamparr['u3'].reshape(-1, 1), clamparr['v3'].reshape(-1, 1)))
    uv4 = np.hstack((clamparr['u4'].reshape(-1, 1), clamparr['v4'].reshape(-1, 1)))
    clamp = clamparr['camp']
    sigma = clamparr['sigmaca']

    # data weighting
    if weighting == 'uniform':
        sigma = np.median(sigma) * np.ones(len(sigma))

    # get NFFT info
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    A1 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv1)
    A2 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv2)
    A3 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv3)
    A4 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv4)
    A = [A1, A2, A3, A4]

    return (clamp, sigma, A)


def chisqdata_logcamp_diag_nfft(Obsdata, Prior, pol='I', **kwargs):
    """Return the diagonalized log closure amplitudes, sigmas, uv points, and nfft info
    """
    if (Prior.xdim % 2 or Prior.ydim % 2):
        raise Exception("NFFT doesn't work with odd image dimensions!")

    # unpack keyword args
    maxset = kwargs.get('maxset', False)
    if maxset:
        count = 'max'
    else:
        count = 'min'

    snrcut = kwargs.get('snrcut', 0.)
    debias = kwargs.get('debias', True)
    fft_pad_factor = kwargs.get('fft_pad_factor', ehc.FFT_PAD_DEFAULT)
    p_rad = kwargs.get('p_rad', ehc.GRIDDER_P_RAD_DEFAULT)

    # unpack data & mask low snr points
    vtype = ehc.vis_poldict[pol]
    clamparr = Obsdata.c_log_amplitudes_diag(vtype=vtype, count=count, debias=debias, snrcut=snrcut)

    # loop over timestamps
    clamp_diag = []
    sigma_diag = []
    tform_mats = []
    u1 = []
    v1 = []
    u2 = []
    v2 = []
    u3 = []
    v3 = []
    u4 = []
    v4 = []
    for ic, cl in enumerate(clamparr):

        # get diagonalized log closure amplitudes and errors
        clamp_diag.append(cl[0]['camp'])
        sigma_diag.append(cl[0]['sigmaca'])

        # get u and v values
        u1.append(cl[2][:, 0].astype('float'))
        v1.append(cl[3][:, 0].astype('float'))
        u2.append(cl[2][:, 1].astype('float'))
        v2.append(cl[3][:, 1].astype('float'))
        u3.append(cl[2][:, 2].astype('float'))
        v3.append(cl[3][:, 2].astype('float'))
        u4.append(cl[2][:, 3].astype('float'))
        v4.append(cl[3][:, 3].astype('float'))

        # get transformation matrix for this timestamp
        tform_mats.append(cl[4].astype('float'))

    # fix formatting of arrays
    u1 = np.concatenate(u1)
    v1 = np.concatenate(v1)
    u2 = np.concatenate(u2)
    v2 = np.concatenate(v2)
    u3 = np.concatenate(u3)
    v3 = np.concatenate(v3)
    u4 = np.concatenate(u4)
    v4 = np.concatenate(v4)

    # get uv arrays
    uv1 = np.hstack((u1.reshape(-1, 1), v1.reshape(-1, 1)))
    uv2 = np.hstack((u2.reshape(-1, 1), v2.reshape(-1, 1)))
    uv3 = np.hstack((u3.reshape(-1, 1), v3.reshape(-1, 1)))
    uv4 = np.hstack((u4.reshape(-1, 1), v4.reshape(-1, 1)))

    # get NFFT info
    npad = int(fft_pad_factor * np.max((Prior.xdim, Prior.ydim)))
    A1 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv1)
    A2 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv2)
    A3 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv3)
    A4 = obsh.NFFTInfo(Prior.xdim, Prior.ydim, Prior.psize, Prior.pulse, npad, p_rad, uv4)
    A = [A1, A2, A3, A4]

    # combine Fourier and transformation matrices into tuple for outputting
    Amatrices = (A, np.array(tform_mats))

    return (np.array(clamp_diag), np.array(sigma_diag), Amatrices)

##################################################################################################
# Restoring ,Embedding, and Plotting Functions
##################################################################################################


def plot_i(im, Prior, nit, chi2_dict, **kwargs):
    """Plot the total intensity image at each iteration
    """

    cmap = kwargs.get('cmap', 'afmhot')
    interpolation = kwargs.get('interpolation', 'gaussian')
    pol = kwargs.get('pol', '')
    scale = kwargs.get('scale', None)
    dynamic_range = kwargs.get('dynamic_range', 1.e5)
    gamma = kwargs.get('dynamic_range', .5)

    plt.ion()
    plt.pause(1.e-6)
    plt.clf()

    imarr = im.reshape(Prior.ydim, Prior.xdim)

    if scale == 'log':
        if (imarr < 0.0).any():
            print('clipping values less than 0')
            imarr[imarr < 0.0] = 0.0
        imarr = np.log(imarr + np.max(imarr)/dynamic_range)

    if scale == 'gamma':
        if (imarr < 0.0).any():
            print('clipping values less than 0')
            imarr[imarr < 0.0] = 0.0
        imarr = (imarr + np.max(imarr)/dynamic_range)**(gamma)

    plt.imshow(imarr, cmap=plt.get_cmap(cmap), interpolation=interpolation)
    xticks = obsh.ticks(Prior.xdim, Prior.psize/ehc.RADPERAS/1e-6)
    yticks = obsh.ticks(Prior.ydim, Prior.psize/ehc.RADPERAS/1e-6)
    plt.xticks(xticks[0], xticks[1])
    plt.yticks(yticks[0], yticks[1])
    plt.xlabel(r'Relative RA ($\mu$as)')
    plt.ylabel(r'Relative Dec ($\mu$as)')
    plotstr = str(pol) + " : step: %i  " % nit
    for key in chi2_dict.keys():
        plotstr += r"$\chi^2_{%s}$: %0.2f  " % (key, chi2_dict[key])
    plt.title(plotstr, fontsize=18)


def embed(im, mask, clipfloor=0., randomfloor=False):
    """Embeds a 1d image array into the size of boolean embed mask
    """

    out = np.zeros(len(mask))

    # Here's a much faster version than before
    out[mask.nonzero()] = im

    if clipfloor != 0.0:
        if randomfloor:  # prevent total variation gradient singularities
            out[(mask-1).nonzero()] = clipfloor * \
                np.abs(np.random.normal(size=len((mask-1).nonzero())))
        else:
            out[(mask-1).nonzero()] = clipfloor

    return out
